<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 (
	"fmt"
	"os"
	"time"
	"bufio"
	"strings"

	"github.com/golang/glog"
	"github.com/spf13/cobra"
	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 create 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
}

// User input routine copied from vendor/google.golang.org/api/examples/gmail.go
func PromptInteractive(upgradedHost string) (stop_prompting bool) {
	stop_prompting = false
	reader := bufio.NewReader(os.Stdin)
	glog.Infof("Pausing after finished %q", upgradedHost)
	fmt.Printf("Continue? (Y)es, (N)o, (A)lwaysYes: [Y] ")
	val := ""
	var err error;
	if val, err = reader.ReadString('\n'); err != nil {
		glog.Fatalf("unable to interpret input: %v", err)
	}
	val = strings.TrimSpace(val)
	val = strings.ToLower(val)
	switch val {
	case "y","","\n":
		glog.V(4).Infof("Continuing with next host (response %q)\n",val)
	case "n":
		glog.Infof("User signaled to stop")
		os.Exit(3)
	case "a":
		glog.Infof("Always Yes, stop prompting for rest of hosts")
		stop_prompting = true
	}
	return stop_prompting
}

// 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, 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, 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, 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 {
				var stop_prompting bool = PromptInteractive(nodeName)
				if stop_prompting {
					// 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, 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, 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, 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, instanceGroupList *api.InstanceGroupList, duration time.Duration, tickDuration time.Duration) bool {
	if _, err := validation.ValidateCluster(rollingUpdateData.ClusterName, 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, instanceGroupList *api.InstanceGroupList) error {

	if _, err := validation.ValidateCluster(rollingUpdateData.ClusterName, instanceGroupList, rollingUpdateData.K8sClient); err != nil {
		return fmt.Errorf("cluster %q did not pass validation: %v", rollingUpdateData.ClusterName, 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.", id, nodeName, r.CloudGroup.HumanName)
	} else {
		glog.Infof("Stopping instance %q, in group %q.", 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,
	}

	cmd := &cobra.Command{
		Use: "cordon NODE",
	}
	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.V(3).Infof("Waiting for %s for pods to stabilize after draining.", rollingUpdateData.PostDrainDelay)
		time.Sleep(rollingUpdateData.PostDrainDelay)
	}

	return nil
}

// Delete and 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>Lint fixes by make gofmt<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"
	"github.com/spf13/cobra"
	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 create 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
}

// User input routine copied from vendor/google.golang.org/api/examples/gmail.go
func PromptInteractive(upgradedHost string) (stop_prompting bool) {
	stop_prompting = false
	reader := bufio.NewReader(os.Stdin)
	glog.Infof("Pausing after finished %q", upgradedHost)
	fmt.Printf("Continue? (Y)es, (N)o, (A)lwaysYes: [Y] ")
	val := ""
	var err error
	if val, err = reader.ReadString('\n'); err != nil {
		glog.Fatalf("unable to interpret input: %v", err)
	}
	val = strings.TrimSpace(val)
	val = strings.ToLower(val)
	switch val {
	case "y", "", "\n":
		glog.V(4).Infof("Continuing with next host (response %q)\n", val)
	case "n":
		glog.Infof("User signaled to stop")
		os.Exit(3)
	case "a":
		glog.Infof("Always Yes, stop prompting for rest of hosts")
		stop_prompting = true
	}
	return stop_prompting
}

// 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, 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, 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, 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 {
				var stop_prompting bool = PromptInteractive(nodeName)
				if stop_prompting {
					// 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, 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, 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, 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, instanceGroupList *api.InstanceGroupList, duration time.Duration, tickDuration time.Duration) bool {
	if _, err := validation.ValidateCluster(rollingUpdateData.ClusterName, 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, instanceGroupList *api.InstanceGroupList) error {

	if _, err := validation.ValidateCluster(rollingUpdateData.ClusterName, instanceGroupList, rollingUpdateData.K8sClient); err != nil {
		return fmt.Errorf("cluster %q did not pass validation: %v", rollingUpdateData.ClusterName, 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.", id, nodeName, r.CloudGroup.HumanName)
	} else {
		glog.Infof("Stopping instance %q, in group %q.", 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,
	}

	cmd := &cobra.Command{
		Use: "cordon NODE",
	}
	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.V(3).Infof("Waiting for %s for pods to stabilize after draining.", rollingUpdateData.PostDrainDelay)
		time.Sleep(rollingUpdateData.PostDrainDelay)
	}

	return nil
}

// Delete and 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>/*
Copyright 2019 The Rook 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 config provides default configurations which Rook will set in Ceph clusters.
package config

import (
	rookceph "github.com/rook/rook/pkg/apis/ceph.rook.io/v1"
	"github.com/rook/rook/pkg/operator/ceph/version"
)

// DefaultFlags returns the default configuration flags Rook will set on the command line for all
// calls to Ceph daemons and tools. Values specified here will not be able to be overridden using
// the mon's central KV store, and that is (and should be) by intent.
func DefaultFlags(fsid, mountedKeyringPath string, cephVersion version.CephVersion) []string {
	flags := []string{
		// fsid unnecessary but is a safety to make sure daemons can only connect to their cluster
		NewFlag("fsid", fsid),
		NewFlag("keyring", mountedKeyringPath),
		// For containers, we're expected to log everything to stderr
		NewFlag("log-to-stderr", "true"),
		NewFlag("err-to-stderr", "true"),
		NewFlag("mon-cluster-log-to-stderr", "true"),
		// differentiate debug text from audit text, and the space after 'debug' is critical
		NewFlag("log-stderr-prefix", "debug "),
	}

	// As of Nautilus 14.2.1 at least
	// These new flags control Ceph's daemon logging behavior to files
	// By default we set them to False so no logs get written on file
	// However they can be activated at any time via the centralized config store
	if cephVersion.IsAtLeast(version.CephVersion{Major: 14, Minor: 2, Extra: 1}) {
		flags = append(flags, []string{
			NewFlag("default-log-to-file", "false"),
			NewFlag("default-mon-cluster-log-to-file", "false"),
		}...)
	}

	flags = append(flags, StoredMonHostEnvVarFlags()...)

	return flags
}

// makes it possible to be slightly less verbose to create a ConfigOverride here
func configOverride(who, option, value string) rookceph.ConfigOverride {
	return rookceph.ConfigOverride{Who: who, Option: option, Value: value}
}

// DefaultCentralizedConfigs returns the default configuration options Rook will set in Ceph's
// centralized config store.
func DefaultCentralizedConfigs(cephVersion version.CephVersion) rookceph.ConfigOverridesSpec {
	overrides := []rookceph.ConfigOverride{
		configOverride("global", "mon allow pool delete", "true"),
	}

	// Everything before Nautilus 14.2.1
	// Prior to Nautilus 14.2.1 certain log flags were not present
	// so in order to not log anything on files we must set the following flags to null
	// Since Nautilus 14.2.1 introduced both 'default-log-to-file' and 'default-mon-cluster-log-to-file' (see above defaultFlagConfigs)
	// these are not needed
	if !cephVersion.IsAtLeast(version.CephVersion{Major: 14, Minor: 2, Extra: 1}) {
		// Set the default log files to empty so they don't bloat containers. Can be changed in
		// Mimic+ by users if needed.
		overrides = append(overrides, []rookceph.ConfigOverride{
			configOverride("global", "log file", ""),
			configOverride("global", "mon cluster log file", ""),
		}...)
	}

	return overrides
}

// DefaultLegacyConfigs need to be added to the Ceph config file until the integration tests can be
// made to override these options for the Ceph clusters it creates.
func DefaultLegacyConfigs() rookceph.ConfigOverridesSpec {
	overrides := []rookceph.ConfigOverride{
		configOverride("global", "mon max pg per osd", "1000"),
		//
		// TODO: remove these; if we need for integration tests, set in integration test spec
		configOverride("global", "osd pool default size", "1"),
		configOverride("global", "osd pool default min size", "1"),
		configOverride("global", "osd pool default pg num", "100"),
		configOverride("global", "osd pool default pgp num", "100"),
		//
		// TODO: drop this when FlexVolume is no longer supported
		configOverride("global", "rbd_default_features", "3"),
	}
	return overrides
}
<commit_msg>Ceph: remove legacy configs<commit_after>/*
Copyright 2019 The Rook 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 config provides default configurations which Rook will set in Ceph clusters.
package config

import (
	rookceph "github.com/rook/rook/pkg/apis/ceph.rook.io/v1"
	"github.com/rook/rook/pkg/operator/ceph/version"
)

// DefaultFlags returns the default configuration flags Rook will set on the command line for all
// calls to Ceph daemons and tools. Values specified here will not be able to be overridden using
// the mon's central KV store, and that is (and should be) by intent.
func DefaultFlags(fsid, mountedKeyringPath string, cephVersion version.CephVersion) []string {
	flags := []string{
		// fsid unnecessary but is a safety to make sure daemons can only connect to their cluster
		NewFlag("fsid", fsid),
		NewFlag("keyring", mountedKeyringPath),
		// For containers, we're expected to log everything to stderr
		NewFlag("log-to-stderr", "true"),
		NewFlag("err-to-stderr", "true"),
		NewFlag("mon-cluster-log-to-stderr", "true"),
		// differentiate debug text from audit text, and the space after 'debug' is critical
		NewFlag("log-stderr-prefix", "debug "),
	}

	// As of Nautilus 14.2.1 at least
	// These new flags control Ceph's daemon logging behavior to files
	// By default we set them to False so no logs get written on file
	// However they can be activated at any time via the centralized config store
	if cephVersion.IsAtLeast(version.CephVersion{Major: 14, Minor: 2, Extra: 1}) {
		flags = append(flags, []string{
			NewFlag("default-log-to-file", "false"),
			NewFlag("default-mon-cluster-log-to-file", "false"),
		}...)
	}

	flags = append(flags, StoredMonHostEnvVarFlags()...)

	return flags
}

// makes it possible to be slightly less verbose to create a ConfigOverride here
func configOverride(who, option, value string) rookceph.ConfigOverride {
	return rookceph.ConfigOverride{Who: who, Option: option, Value: value}
}

// DefaultCentralizedConfigs returns the default configuration options Rook will set in Ceph's
// centralized config store.
func DefaultCentralizedConfigs(cephVersion version.CephVersion) rookceph.ConfigOverridesSpec {
	overrides := []rookceph.ConfigOverride{
		configOverride("global", "mon allow pool delete", "true"),
	}

	// Everything before Nautilus 14.2.1
	// Prior to Nautilus 14.2.1 certain log flags were not present
	// so in order to not log anything on files we must set the following flags to null
	// Since Nautilus 14.2.1 introduced both 'default-log-to-file' and 'default-mon-cluster-log-to-file' (see above defaultFlagConfigs)
	// these are not needed
	if !cephVersion.IsAtLeast(version.CephVersion{Major: 14, Minor: 2, Extra: 1}) {
		// Set the default log files to empty so they don't bloat containers. Can be changed in
		// Mimic+ by users if needed.
		overrides = append(overrides, []rookceph.ConfigOverride{
			configOverride("global", "log file", ""),
			configOverride("global", "mon cluster log file", ""),
		}...)
	}

	return overrides
}

// DefaultLegacyConfigs need to be added to the Ceph config file until the integration tests can be
// made to override these options for the Ceph clusters it creates.
func DefaultLegacyConfigs() rookceph.ConfigOverridesSpec {
	overrides := []rookceph.ConfigOverride{
		// TODO: drop this when FlexVolume is no longer supported
		configOverride("global", "rbd_default_features", "3"),
	}
	return overrides
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Rook 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 pool to manage a rook pool.
package pool

import (
	"context"

	"github.com/coreos/pkg/capnslog"
	cephclient "github.com/rook/rook/pkg/daemon/ceph/client"

	"github.com/pkg/errors"
	cephv1 "github.com/rook/rook/pkg/apis/ceph.rook.io/v1"
	"github.com/rook/rook/pkg/clusterd"
	"github.com/rook/rook/pkg/daemon/ceph/model"
	opcontroller "github.com/rook/rook/pkg/operator/ceph/controller"
	"github.com/rook/rook/pkg/operator/k8sutil"
	kerrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/handler"
	"sigs.k8s.io/controller-runtime/pkg/manager"
	"sigs.k8s.io/controller-runtime/pkg/reconcile"
	"sigs.k8s.io/controller-runtime/pkg/source"
)

const (
	replicatedType         = "replicated"
	erasureCodeType        = "erasure-coded"
	poolApplicationNameRBD = "rbd"
	controllerName         = "ceph-block-pool-controller"
)

var logger = capnslog.NewPackageLogger("github.com/rook/rook", controllerName)

var _ reconcile.Reconciler = &ReconcileCephBlockPool{}

// ReconcileCephBlockPool reconciles a CephBlockPool object
type ReconcileCephBlockPool struct {
	client  client.Client
	scheme  *runtime.Scheme
	context *clusterd.Context
}

// Add creates a new CephBlockPool Controller and adds it to the Manager. The Manager will set fields on the Controller
// and Start it when the Manager is Started.
func Add(mgr manager.Manager, context *clusterd.Context) error {
	return add(mgr, newReconciler(mgr, context))
}

// newReconciler returns a new reconcile.Reconciler
func newReconciler(mgr manager.Manager, context *clusterd.Context) reconcile.Reconciler {
	// Add the cephv1 scheme to the manager scheme so that the controller knows about it
	mgrScheme := mgr.GetScheme()
	cephv1.AddToScheme(mgr.GetScheme())

	return &ReconcileCephBlockPool{
		client:  mgr.GetClient(),
		scheme:  mgrScheme,
		context: context,
	}
}

func add(mgr manager.Manager, r reconcile.Reconciler) error {
	// Create a new controller
	c, err := controller.New(controllerName, mgr, controller.Options{Reconciler: r})
	if err != nil {
		return err
	}

	// Watch for changes on the CephBlockPool CRD object
	err = c.Watch(&source.Kind{Type: &cephv1.CephBlockPool{}}, &handler.EnqueueRequestForObject{}, opcontroller.WatchUpdatePredicate())
	if err != nil {
		return err
	}

	return nil
}

// Reconcile reads that state of the cluster for a CephBlockPool object and makes changes based on the state read
// and what is in the CephBlockPool.Spec
// The Controller will requeue the Request to be processed again if the returned error is non-nil or
// Result.Requeue is true, otherwise upon completion it will remove the work from the queue.
func (r *ReconcileCephBlockPool) Reconcile(request reconcile.Request) (reconcile.Result, error) {
	// workaround because the rook logging mechanism is not compatible with the controller-runtime loggin interface
	reconcileResponse, err := r.reconcile(request)
	if err != nil {
		logger.Errorf("failed to reconcile %v", err)
	}

	return reconcileResponse, err
}

func (r *ReconcileCephBlockPool) reconcile(request reconcile.Request) (reconcile.Result, error) {
	// Fetch the CephBlockPool instance
	cephBlockPool := &cephv1.CephBlockPool{}
	err := r.client.Get(context.TODO(), request.NamespacedName, cephBlockPool)
	if err != nil {
		if kerrors.IsNotFound(err) {
			logger.Debug("CephBlockPool resource not found. Ignoring since object must be deleted.")
			return reconcile.Result{}, nil
		}
		// Error reading the object - requeue the request.
		return reconcile.Result{}, errors.Wrapf(err, "failed to get CephBlockPool")
	}

	// The CR was just created, initializing status fields
	if cephBlockPool.Status == nil {
		cephBlockPool.Status = &cephv1.Status{}
		cephBlockPool.Status.Phase = k8sutil.Created
		err := opcontroller.UpdateStatus(r.client, cephBlockPool)
		if err != nil {
			return reconcile.Result{}, errors.Wrap(err, "failed to set status")
		}
	}

	// Make sure a CephCluster is present otherwise do nothing
	_, isReadyToReconcile, cephClusterExists, reconcileResponse := opcontroller.IsReadyToReconcile(r.client, r.context, request.NamespacedName)
	if !isReadyToReconcile {
		// This handles the case where the Ceph Cluster is gone and we want to delete that CR
		// We skip the deletePool() function since everything is gone already
		//
		// ALso, only remove the finalizer if the CephCluster is gone
		// If not, we should wait for it to be ready
		// This handles the case where the operator is not ready to accept Ceph command but the cluster exists
		if !cephBlockPool.GetDeletionTimestamp().IsZero() && !cephClusterExists {
			// Remove finalizer
			err = opcontroller.RemoveFinalizer(r.client, cephBlockPool)
			if err != nil {
				return reconcile.Result{}, errors.Wrap(err, "failed to remove finalizer")
			}

			// Return and do not requeue. Successful deletion.
			return reconcile.Result{}, nil
		}

		logger.Debugf("CephCluster resource not ready in namespace %q, retrying in %q.", request.NamespacedName.Namespace, opcontroller.WaitForRequeueIfCephClusterNotReadyAfter.String())
		return reconcileResponse, nil
	}

	// Set a finalizer so we can do cleanup before the object goes away
	err = opcontroller.AddFinalizerIfNotPresent(r.client, cephBlockPool)
	if err != nil {
		return reconcile.Result{}, errors.Wrap(err, "failed to add finalizer")
	}

	// DELETE: the CR was deleted
	if !cephBlockPool.GetDeletionTimestamp().IsZero() {
		logger.Debugf("deleting pool %q", cephBlockPool.Name)
		err := deletePool(r.context, cephBlockPool)
		if err != nil {
			return reconcile.Result{}, errors.Wrapf(err, "failed to delete pool %q. ", cephBlockPool.Name)
		}

		// Remove finalizer
		err = opcontroller.RemoveFinalizer(r.client, cephBlockPool)
		if err != nil {
			return reconcile.Result{}, errors.Wrap(err, "failed to remove finalizer")
		}

		// Return and do not requeue. Successful deletion.
		return reconcile.Result{}, nil
	}

	// validate the pool settings
	if err := ValidatePool(r.context, cephBlockPool); err != nil {
		return reconcile.Result{}, errors.Wrapf(err, "invalid pool CR %q spec", cephBlockPool.Name)
	}

	// Start object reconciliation, updating status for this
	cephBlockPool.Status.Phase = k8sutil.ReconcilingStatus
	err = opcontroller.UpdateStatus(r.client, cephBlockPool)
	if err != nil {
		return reconcile.Result{}, errors.Wrap(err, "failed to set status")
	}

	// CREATE/UPDATE
	reconcileResponse, err = r.reconcileCreatePool(cephBlockPool)
	if err != nil {
		cephBlockPool.Status.Phase = k8sutil.ReconcileFailedStatus
		errStatus := opcontroller.UpdateStatus(r.client, cephBlockPool)
		if errStatus != nil {
			return reconcile.Result{}, errors.Wrap(errStatus, "failed to set status")
		}
		return reconcileResponse, errors.Wrapf(err, "failed to create pool %q.", cephBlockPool.GetName())
	}

	// Set Ready status, we are done reconciling
	cephBlockPool.Status.Phase = k8sutil.ReadyStatus
	err = opcontroller.UpdateStatus(r.client, cephBlockPool)
	if err != nil {
		return reconcile.Result{}, errors.Wrap(err, "failed to set status")
	}

	// Return and do not requeue
	logger.Debug("done reconciling")
	return reconcile.Result{}, nil
}

func (r *ReconcileCephBlockPool) reconcileCreatePool(cephBlockPool *cephv1.CephBlockPool) (reconcile.Result, error) {
	err := createPool(r.context, cephBlockPool)
	if err != nil {
		return reconcile.Result{}, errors.Wrapf(err, "failed to create pool %q.", cephBlockPool.GetName())
	}

	// Let's return here so that on the initial creation we don't check for update right away
	return reconcile.Result{}, nil
}

// Create the pool
func createPool(context *clusterd.Context, p *cephv1.CephBlockPool) error {
	// create the pool
	logger.Infof("creating pool %q in namespace %q", p.Name, p.Namespace)
	if err := cephclient.CreatePoolWithProfile(context, p.Namespace, *p.Spec.ToModel(p.Name), poolApplicationNameRBD); err != nil {
		return errors.Wrapf(err, "failed to create pool %q", p.Name)
	}

	return nil
}

// Delete the pool
func deletePool(context *clusterd.Context, p *cephv1.CephBlockPool) error {
	if err := cephclient.DeletePool(context, p.Namespace, p.Name); err != nil {
		return errors.Wrapf(err, "failed to delete pool %q", p.Name)
	}

	return nil
}

// ModelToSpec reflect the internal pool struct from a pool spec
func ModelToSpec(pool model.Pool) cephv1.PoolSpec {
	ec := pool.ErasureCodedConfig
	return cephv1.PoolSpec{
		FailureDomain: pool.FailureDomain,
		CrushRoot:     pool.CrushRoot,
		DeviceClass:   pool.DeviceClass,
		Replicated:    cephv1.ReplicatedSpec{Size: pool.ReplicatedConfig.Size},
		ErasureCoded:  cephv1.ErasureCodedSpec{CodingChunks: ec.CodingChunkCount, DataChunks: ec.DataChunkCount, Algorithm: ec.Algorithm},
	}
}

// ValidatePool Validate the pool arguments
func ValidatePool(context *clusterd.Context, p *cephv1.CephBlockPool) error {
	if p.Name == "" {
		return errors.New("missing name")
	}
	if p.Namespace == "" {
		return errors.New("missing namespace")
	}
	if err := ValidatePoolSpec(context, p.Namespace, &p.Spec); err != nil {
		return err
	}
	return nil
}

// ValidatePoolSpec validates the Ceph block pool spec CR
func ValidatePoolSpec(context *clusterd.Context, namespace string, p *cephv1.PoolSpec) error {
	if p.Replication() != nil && p.ErasureCode() != nil {
		return errors.New("both replication and erasure code settings cannot be specified")
	}

	var crush cephclient.CrushMap
	var err error
	if p.FailureDomain != "" || p.CrushRoot != "" {
		crush, err = cephclient.GetCrushMap(context, namespace)
		if err != nil {
			return errors.Wrapf(err, "failed to get crush map")
		}
	}

	// validate the failure domain if specified
	if p.FailureDomain != "" {
		found := false
		for _, t := range crush.Types {
			if t.Name == p.FailureDomain {
				found = true
				break
			}
		}
		if !found {
			return errors.Errorf("unrecognized failure domain %s", p.FailureDomain)
		}
	}

	// validate the crush root if specified
	if p.CrushRoot != "" {
		found := false
		for _, t := range crush.Buckets {
			if t.Name == p.CrushRoot {
				found = true
				break
			}
		}
		if !found {
			return errors.Errorf("unrecognized crush root %s", p.CrushRoot)
		}
	}

	// validate pool replica size
	if p.Replicated.Size == 1 && p.Replicated.RequireSafeReplicaSize {
		return errors.Errorf("error pool size is %d and requireSafeReplicaSize is %t, must be false", p.Replicated.Size, p.Replicated.RequireSafeReplicaSize)
	}

	return nil
}
<commit_msg>ceph: only delete pool if exists<commit_after>/*
Copyright 2016 The Rook 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 pool to manage a rook pool.
package pool

import (
	"context"

	"github.com/coreos/pkg/capnslog"
	cephclient "github.com/rook/rook/pkg/daemon/ceph/client"

	"github.com/pkg/errors"
	cephv1 "github.com/rook/rook/pkg/apis/ceph.rook.io/v1"
	"github.com/rook/rook/pkg/clusterd"
	"github.com/rook/rook/pkg/daemon/ceph/model"
	opcontroller "github.com/rook/rook/pkg/operator/ceph/controller"
	"github.com/rook/rook/pkg/operator/k8sutil"
	kerrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/handler"
	"sigs.k8s.io/controller-runtime/pkg/manager"
	"sigs.k8s.io/controller-runtime/pkg/reconcile"
	"sigs.k8s.io/controller-runtime/pkg/source"
)

const (
	replicatedType         = "replicated"
	erasureCodeType        = "erasure-coded"
	poolApplicationNameRBD = "rbd"
	controllerName         = "ceph-block-pool-controller"
)

var logger = capnslog.NewPackageLogger("github.com/rook/rook", controllerName)

var _ reconcile.Reconciler = &ReconcileCephBlockPool{}

// ReconcileCephBlockPool reconciles a CephBlockPool object
type ReconcileCephBlockPool struct {
	client  client.Client
	scheme  *runtime.Scheme
	context *clusterd.Context
}

// Add creates a new CephBlockPool Controller and adds it to the Manager. The Manager will set fields on the Controller
// and Start it when the Manager is Started.
func Add(mgr manager.Manager, context *clusterd.Context) error {
	return add(mgr, newReconciler(mgr, context))
}

// newReconciler returns a new reconcile.Reconciler
func newReconciler(mgr manager.Manager, context *clusterd.Context) reconcile.Reconciler {
	// Add the cephv1 scheme to the manager scheme so that the controller knows about it
	mgrScheme := mgr.GetScheme()
	cephv1.AddToScheme(mgr.GetScheme())

	return &ReconcileCephBlockPool{
		client:  mgr.GetClient(),
		scheme:  mgrScheme,
		context: context,
	}
}

func add(mgr manager.Manager, r reconcile.Reconciler) error {
	// Create a new controller
	c, err := controller.New(controllerName, mgr, controller.Options{Reconciler: r})
	if err != nil {
		return err
	}

	// Watch for changes on the CephBlockPool CRD object
	err = c.Watch(&source.Kind{Type: &cephv1.CephBlockPool{}}, &handler.EnqueueRequestForObject{}, opcontroller.WatchUpdatePredicate())
	if err != nil {
		return err
	}

	return nil
}

// Reconcile reads that state of the cluster for a CephBlockPool object and makes changes based on the state read
// and what is in the CephBlockPool.Spec
// The Controller will requeue the Request to be processed again if the returned error is non-nil or
// Result.Requeue is true, otherwise upon completion it will remove the work from the queue.
func (r *ReconcileCephBlockPool) Reconcile(request reconcile.Request) (reconcile.Result, error) {
	// workaround because the rook logging mechanism is not compatible with the controller-runtime loggin interface
	reconcileResponse, err := r.reconcile(request)
	if err != nil {
		logger.Errorf("failed to reconcile %v", err)
	}

	return reconcileResponse, err
}

func (r *ReconcileCephBlockPool) reconcile(request reconcile.Request) (reconcile.Result, error) {
	// Fetch the CephBlockPool instance
	cephBlockPool := &cephv1.CephBlockPool{}
	err := r.client.Get(context.TODO(), request.NamespacedName, cephBlockPool)
	if err != nil {
		if kerrors.IsNotFound(err) {
			logger.Debug("CephBlockPool resource not found. Ignoring since object must be deleted.")
			return reconcile.Result{}, nil
		}
		// Error reading the object - requeue the request.
		return reconcile.Result{}, errors.Wrapf(err, "failed to get CephBlockPool")
	}

	// The CR was just created, initializing status fields
	if cephBlockPool.Status == nil {
		cephBlockPool.Status = &cephv1.Status{}
		cephBlockPool.Status.Phase = k8sutil.Created
		err := opcontroller.UpdateStatus(r.client, cephBlockPool)
		if err != nil {
			return reconcile.Result{}, errors.Wrap(err, "failed to set status")
		}
	}

	// Make sure a CephCluster is present otherwise do nothing
	_, isReadyToReconcile, cephClusterExists, reconcileResponse := opcontroller.IsReadyToReconcile(r.client, r.context, request.NamespacedName)
	if !isReadyToReconcile {
		// This handles the case where the Ceph Cluster is gone and we want to delete that CR
		// We skip the deletePool() function since everything is gone already
		//
		// ALso, only remove the finalizer if the CephCluster is gone
		// If not, we should wait for it to be ready
		// This handles the case where the operator is not ready to accept Ceph command but the cluster exists
		if !cephBlockPool.GetDeletionTimestamp().IsZero() && !cephClusterExists {
			// Remove finalizer
			err = opcontroller.RemoveFinalizer(r.client, cephBlockPool)
			if err != nil {
				return reconcile.Result{}, errors.Wrap(err, "failed to remove finalizer")
			}

			// Return and do not requeue. Successful deletion.
			return reconcile.Result{}, nil
		}

		logger.Debugf("CephCluster resource not ready in namespace %q, retrying in %q.", request.NamespacedName.Namespace, opcontroller.WaitForRequeueIfCephClusterNotReadyAfter.String())
		return reconcileResponse, nil
	}

	// Set a finalizer so we can do cleanup before the object goes away
	err = opcontroller.AddFinalizerIfNotPresent(r.client, cephBlockPool)
	if err != nil {
		return reconcile.Result{}, errors.Wrap(err, "failed to add finalizer")
	}

	// DELETE: the CR was deleted
	if !cephBlockPool.GetDeletionTimestamp().IsZero() {
		logger.Debugf("deleting pool %q", cephBlockPool.Name)
		err := deletePool(r.context, cephBlockPool)
		if err != nil {
			return reconcile.Result{}, errors.Wrapf(err, "failed to delete pool %q. ", cephBlockPool.Name)
		}

		// Remove finalizer
		err = opcontroller.RemoveFinalizer(r.client, cephBlockPool)
		if err != nil {
			return reconcile.Result{}, errors.Wrap(err, "failed to remove finalizer")
		}

		// Return and do not requeue. Successful deletion.
		return reconcile.Result{}, nil
	}

	// validate the pool settings
	if err := ValidatePool(r.context, cephBlockPool); err != nil {
		return reconcile.Result{}, errors.Wrapf(err, "invalid pool CR %q spec", cephBlockPool.Name)
	}

	// Start object reconciliation, updating status for this
	cephBlockPool.Status.Phase = k8sutil.ReconcilingStatus
	err = opcontroller.UpdateStatus(r.client, cephBlockPool)
	if err != nil {
		return reconcile.Result{}, errors.Wrap(err, "failed to set status")
	}

	// CREATE/UPDATE
	reconcileResponse, err = r.reconcileCreatePool(cephBlockPool)
	if err != nil {
		cephBlockPool.Status.Phase = k8sutil.ReconcileFailedStatus
		errStatus := opcontroller.UpdateStatus(r.client, cephBlockPool)
		if errStatus != nil {
			return reconcile.Result{}, errors.Wrap(errStatus, "failed to set status")
		}
		return reconcileResponse, errors.Wrapf(err, "failed to create pool %q.", cephBlockPool.GetName())
	}

	// Set Ready status, we are done reconciling
	cephBlockPool.Status.Phase = k8sutil.ReadyStatus
	err = opcontroller.UpdateStatus(r.client, cephBlockPool)
	if err != nil {
		return reconcile.Result{}, errors.Wrap(err, "failed to set status")
	}

	// Return and do not requeue
	logger.Debug("done reconciling")
	return reconcile.Result{}, nil
}

func (r *ReconcileCephBlockPool) reconcileCreatePool(cephBlockPool *cephv1.CephBlockPool) (reconcile.Result, error) {
	err := createPool(r.context, cephBlockPool)
	if err != nil {
		return reconcile.Result{}, errors.Wrapf(err, "failed to create pool %q.", cephBlockPool.GetName())
	}

	// Let's return here so that on the initial creation we don't check for update right away
	return reconcile.Result{}, nil
}

// Create the pool
func createPool(context *clusterd.Context, p *cephv1.CephBlockPool) error {
	// create the pool
	logger.Infof("creating pool %q in namespace %q", p.Name, p.Namespace)
	if err := cephclient.CreatePoolWithProfile(context, p.Namespace, *p.Spec.ToModel(p.Name), poolApplicationNameRBD); err != nil {
		return errors.Wrapf(err, "failed to create pool %q", p.Name)
	}

	return nil
}

// Delete the pool
func deletePool(context *clusterd.Context, p *cephv1.CephBlockPool) error {
	pools, err := cephclient.ListPoolSummaries(context, p.Namespace)
	if err != nil {
		return errors.Wrapf(err, "failed to list pools")
	}

	// Only delete the pool if it exists...
	for _, pool := range pools {
		if pool.Name == p.Name {
			err := cephclient.DeletePool(context, p.Namespace, p.Name)
			if err != nil {
				return errors.Wrapf(err, "failed to delete pool %q", p.Name)
			}
		}
	}

	return nil
}

// ModelToSpec reflect the internal pool struct from a pool spec
func ModelToSpec(pool model.Pool) cephv1.PoolSpec {
	ec := pool.ErasureCodedConfig
	return cephv1.PoolSpec{
		FailureDomain: pool.FailureDomain,
		CrushRoot:     pool.CrushRoot,
		DeviceClass:   pool.DeviceClass,
		Replicated:    cephv1.ReplicatedSpec{Size: pool.ReplicatedConfig.Size},
		ErasureCoded:  cephv1.ErasureCodedSpec{CodingChunks: ec.CodingChunkCount, DataChunks: ec.DataChunkCount, Algorithm: ec.Algorithm},
	}
}

// ValidatePool Validate the pool arguments
func ValidatePool(context *clusterd.Context, p *cephv1.CephBlockPool) error {
	if p.Name == "" {
		return errors.New("missing name")
	}
	if p.Namespace == "" {
		return errors.New("missing namespace")
	}
	if err := ValidatePoolSpec(context, p.Namespace, &p.Spec); err != nil {
		return err
	}
	return nil
}

// ValidatePoolSpec validates the Ceph block pool spec CR
func ValidatePoolSpec(context *clusterd.Context, namespace string, p *cephv1.PoolSpec) error {
	if p.Replication() != nil && p.ErasureCode() != nil {
		return errors.New("both replication and erasure code settings cannot be specified")
	}

	var crush cephclient.CrushMap
	var err error
	if p.FailureDomain != "" || p.CrushRoot != "" {
		crush, err = cephclient.GetCrushMap(context, namespace)
		if err != nil {
			return errors.Wrapf(err, "failed to get crush map")
		}
	}

	// validate the failure domain if specified
	if p.FailureDomain != "" {
		found := false
		for _, t := range crush.Types {
			if t.Name == p.FailureDomain {
				found = true
				break
			}
		}
		if !found {
			return errors.Errorf("unrecognized failure domain %s", p.FailureDomain)
		}
	}

	// validate the crush root if specified
	if p.CrushRoot != "" {
		found := false
		for _, t := range crush.Buckets {
			if t.Name == p.CrushRoot {
				found = true
				break
			}
		}
		if !found {
			return errors.Errorf("unrecognized crush root %s", p.CrushRoot)
		}
	}

	// validate pool replica size
	if p.Replicated.Size == 1 && p.Replicated.RequireSafeReplicaSize {
		return errors.Errorf("error pool size is %d and requireSafeReplicaSize is %t, must be false", p.Replicated.Size, p.Replicated.RequireSafeReplicaSize)
	}

	return nil
}
<|endoftext|>


<commit_before>package pingone

import (
	"bytes"
	"io/ioutil"
	"testing"

	"github.com/PuerkitoBio/goquery"
	"github.com/stretchr/testify/require"
)

var docTests = []struct {
	fn       func(*goquery.Document) bool
	file     string
	expected bool
}{
	{docIsFormSelectDevice, "example/selectdevice.html", true},
}

func TestDocTypes(t *testing.T) {
	for _, tt := range docTests {
		data, err := ioutil.ReadFile(tt.file)
		require.Nil(t, err)

		doc, err := goquery.NewDocumentFromReader(bytes.NewReader(data))
		require.Nil(t, err)

		if tt.fn(doc) != tt.expected {
			t.Errorf("expect doc check of %v to be %v", tt.file, tt.expected)
		}
	}
}
<commit_msg>add test for makeAbsoluteURL method<commit_after>package pingone

import (
	"bytes"
	"io/ioutil"
	"testing"

	"github.com/PuerkitoBio/goquery"
	"github.com/stretchr/testify/require"
)

var docTests = []struct {
	fn       func(*goquery.Document) bool
	file     string
	expected bool
}{
	{docIsFormSelectDevice, "example/selectdevice.html", true},
}

func TestMakeAbsoluteURL(t *testing.T) {
	require.Equal(t, makeAbsoluteURL("/pingid/ppm/devices", "https://authentication.pingone.com"), "https://authentication.pingone.com/pingid/ppm/devices")
	require.Equal(t, makeAbsoluteURL("/pingid/ppm/devices", "https://authentication.pingone.com/"), "https://authentication.pingone.com/pingid/ppm/devices")
}

func TestDocTypes(t *testing.T) {
	for _, tt := range docTests {
		data, err := ioutil.ReadFile(tt.file)
		require.Nil(t, err)

		doc, err := goquery.NewDocumentFromReader(bytes.NewReader(data))
		require.Nil(t, err)

		if tt.fn(doc) != tt.expected {
			t.Errorf("expect doc check of %v to be %v", tt.file, tt.expected)
		}
	}
}
<|endoftext|>


<commit_before>package migrations

import (
	"fmt"
	"io"
	"os"
	"runtime"
	"strings"

	"github.com/cozy/cozy-stack/pkg/config"
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/instance"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/vfs/vfsswift"
	"github.com/cozy/swift"
	multierror "github.com/hashicorp/go-multierror"
)

const swiftV1ContainerPrefixCozy = "cozy-"
const swiftV1ContainerPrefixData = "data-"
const swiftV2ContainerPrefixCozy = "cozy-v2-"
const swiftV2ContainerPrefixData = "data-v2"
const versionSuffix = "-version"
const dirContentType = "directory"

func init() {
	jobs.AddWorker(&jobs.WorkerConfig{
		WorkerType:   "migrations",
		Concurrency:  runtime.NumCPU(),
		MaxExecCount: 2,
		WorkerFunc:   worker,
		WorkerCommit: commit,
	})
}

const swiftV1ToV2 = "swift-v1-to-v2"

type message struct {
	Type    string `json:"type"`
	Cluster int    `json:"cluster"`
}

func worker(ctx *jobs.WorkerContext) error {
	domain := ctx.Domain()

	var msg message
	if err := ctx.UnmarshalMessage(&msg); err != nil {
		return err
	}

	switch msg.Type {
	case swiftV1ToV2:
		return migrateSwiftV1ToV2(domain)
	default:
		return fmt.Errorf("unknown migration type %q", msg.Type)
	}
}

func commit(ctx *jobs.WorkerContext, err error) error {
	if err != nil {
		return nil
	}
	domain := ctx.Domain()

	var msg message
	if err := ctx.UnmarshalMessage(&msg); err != nil {
		return err
	}

	switch msg.Type {
	case swiftV1ToV2:
		return commitSwiftV1ToV2(domain, msg.Cluster)
	default:
		return fmt.Errorf("unknown migration type %q", msg.Type)
	}
}

type object struct {
	obj          swift.Object
	containerSrc string
	containerDst string
}

func migrateSwiftV1ToV2(domain string) error {
	c := config.GetSwiftConnection()
	inst, err := instance.Get(domain)
	if err != nil {
		return err
	}
	if inst.SwiftCluster > 0 {
		return nil
	}

	containerV1 := swiftV1ContainerPrefixCozy + domain
	containerV2 := swiftV2ContainerPrefixCozy + domain

	// container containing thumbnails
	containerV1Data := swiftV1ContainerPrefixData + domain
	containerV2Data := swiftV2ContainerPrefixData + domain

	err = c.VersionContainerCreate(containerV2, containerV2+versionSuffix)
	if err != nil {
		return err
	}

	objc := make(chan object)
	errc := make(chan error)

	go func() {
		errc <- readObjects(c, objc, containerV1, containerV2)
		errc <- readObjects(c, objc, containerV1Data, containerV2Data)
		close(objc)
	}()

	const N = 4

	for i := 0; i < N; i++ {
		go copyObjects(c, inst, objc, errc)
	}

	var errm error
	done := N
	for {
		err := <-errc
		if err != nil {
			errm = multierror.Append(errm, err)
		} else {
			done--
		}
		if done == 0 {
			break
		}
	}

	return errm
}

func commitSwiftV1ToV2(domain string, swiftCluster int) error {
	inst, err := instance.Get(domain)
	if err != nil {
		return err
	}
	if swiftCluster == 0 {
		swiftCluster = 1
	}
	inst.SwiftCluster = swiftCluster
	return instance.Update(inst)
}

func readObjects(c *swift.Connection, objc chan object,
	containerSrc, containerDst string) error {
	return c.ObjectsWalk(containerSrc, nil, func(opts *swift.ObjectsOpts) (interface{}, error) {
		objs, err := c.Objects(containerSrc, opts)
		if err != nil {
			return nil, err
		}
		for _, obj := range objs {
			objc <- object{
				obj:          obj,
				containerSrc: containerSrc,
				containerDst: containerDst,
			}
		}
		return objs, err
	})
}

func copyObjects(c *swift.Connection, db couchdb.Database,
	objc chan object,
	errc chan error) {

	copyBuffer := make([]byte, 128*1024)

	for obj := range objc {
		var err error
		containerSrc := obj.containerSrc
		containerDst := obj.containerDst
		switch {
		case strings.HasPrefix(containerSrc, swiftV1ContainerPrefixCozy):
			err = copyFileDataObject(c, db, containerSrc, containerDst, obj.obj, copyBuffer)
		case strings.HasPrefix(containerSrc, swiftV1ContainerPrefixData):
			err = copyThumbnailDataObject(c, db, containerSrc, containerDst, obj.obj, copyBuffer)
		}
		if err != nil {
			errc <- err
		}
	}

	errc <- nil
}

func copyFileDataObject(c *swift.Connection, db couchdb.Database,
	containerSrc, containerDst string,
	objSrc swift.Object,
	copyBuffer []byte) error {
	if objSrc.ContentType == dirContentType {
		return nil
	}
	dirID, name, ok := splitV2ObjectName(objSrc.Name)
	if !ok {
		return nil
	}
	var res couchdb.ViewResponse
	err := couchdb.ExecView(db, consts.FilesByParentView, &couchdb.ViewRequest{
		Key:         []string{dirID, consts.FileType, name},
		IncludeDocs: false,
	}, &res)
	if err != nil {
		return err
	}
	if len(res.Rows) == 0 {
		return os.ErrNotExist
	}
	objNameDst := vfsswift.MakeObjectName(res.Rows[0].ID)
	return copyObject(c, db, containerSrc, containerDst, objSrc, objNameDst, copyBuffer)
}

func copyThumbnailDataObject(c *swift.Connection, db couchdb.Database,
	containerSrc, containerDst string,
	objSrc swift.Object,
	copyBuffer []byte) error {

	split := strings.SplitN(strings.TrimPrefix(objSrc.Name, "thumbs/"), "-", 2)
	if len(split) != 2 {
		return nil
	}
	objNameDst := "thumbs/" + vfsswift.MakeObjectName(split[0]) + "-" + split[1]
	return copyObject(c, db, containerSrc, containerDst, objSrc, objNameDst, copyBuffer)
}

func copyObject(c *swift.Connection, db couchdb.Database,
	containerSrc, containerDst string,
	objSrc swift.Object, objNameDst string,
	copyBuffer []byte) (err error) {
	infosDst, _, err := c.Object(containerDst, objNameDst)
	if err == nil && infosDst.Hash == objSrc.Hash {
		return nil
	}
	if err != swift.ObjectNotFound {
		return err
	}

	src, h, err := c.ObjectOpen(containerSrc, objSrc.Name, false, nil)
	if err != nil {
		return err
	}
	defer src.Close()

	dst, err := c.ObjectCreate(containerDst, objNameDst, true, objSrc.Hash, objSrc.ContentType, h)
	if err != nil {
		return err
	}
	defer func() {
		if errc := dst.Close(); errc != nil && err == nil {
			err = errc
		}
	}()

	_, err = io.CopyBuffer(dst, src, copyBuffer)
	return err
}

func splitV2ObjectName(objName string) (dirID string, name string, ok bool) {
	split := strings.SplitN(objName, "/", 2)
	if len(split) != 2 {
		return
	}
	dirID, name = split[0], split[1]
	ok = true
	return
}
<commit_msg>Set metadata on migrated containers<commit_after>package migrations

import (
	"fmt"
	"io"
	"os"
	"runtime"
	"strings"

	"github.com/cozy/cozy-stack/pkg/config"
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/instance"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/vfs/vfsswift"
	"github.com/cozy/swift"
	multierror "github.com/hashicorp/go-multierror"
)

const swiftV1ContainerPrefixCozy = "cozy-"
const swiftV1ContainerPrefixData = "data-"
const swiftV2ContainerPrefixCozy = "cozy-v2-"
const swiftV2ContainerPrefixData = "data-v2"
const versionSuffix = "-version"
const dirContentType = "directory"

func init() {
	jobs.AddWorker(&jobs.WorkerConfig{
		WorkerType:   "migrations",
		Concurrency:  runtime.NumCPU(),
		MaxExecCount: 2,
		WorkerFunc:   worker,
		WorkerCommit: commit,
	})
}

const swiftV1ToV2 = "swift-v1-to-v2"

type message struct {
	Type    string `json:"type"`
	Cluster int    `json:"cluster"`
}

func worker(ctx *jobs.WorkerContext) error {
	domain := ctx.Domain()

	var msg message
	if err := ctx.UnmarshalMessage(&msg); err != nil {
		return err
	}

	switch msg.Type {
	case swiftV1ToV2:
		return migrateSwiftV1ToV2(domain)
	default:
		return fmt.Errorf("unknown migration type %q", msg.Type)
	}
}

func commit(ctx *jobs.WorkerContext, err error) error {
	if err != nil {
		return nil
	}
	domain := ctx.Domain()

	var msg message
	if err := ctx.UnmarshalMessage(&msg); err != nil {
		return err
	}

	switch msg.Type {
	case swiftV1ToV2:
		return commitSwiftV1ToV2(domain, msg.Cluster)
	default:
		return fmt.Errorf("unknown migration type %q", msg.Type)
	}
}

type object struct {
	obj          swift.Object
	containerSrc string
	containerDst string
}

func migrateSwiftV1ToV2(domain string) error {
	c := config.GetSwiftConnection()
	inst, err := instance.Get(domain)
	if err != nil {
		return err
	}
	if inst.SwiftCluster > 0 {
		return nil
	}

	containerV1 := swiftV1ContainerPrefixCozy + domain
	containerV2 := swiftV2ContainerPrefixCozy + domain

	// container containing thumbnails
	containerV1Data := swiftV1ContainerPrefixData + domain
	containerV2Data := swiftV2ContainerPrefixData + domain

	err = c.VersionContainerCreate(containerV2, containerV2+versionSuffix)
	if err != nil {
		return err
	}

	objc := make(chan object)
	errc := make(chan error)

	go func() {
		errc <- readObjects(c, objc, containerV1, containerV2)
		errc <- readObjects(c, objc, containerV1Data, containerV2Data)
		close(objc)
	}()

	const N = 4

	for i := 0; i < N; i++ {
		go copyObjects(c, inst, objc, errc)
	}

	var errm error
	done := N
	for {
		err := <-errc
		if err != nil {
			errm = multierror.Append(errm, err)
		} else {
			done--
		}
		if done == 0 {
			break
		}
	}

	return errm
}

func commitSwiftV1ToV2(domain string, swiftCluster int) error {
	inst, err := instance.Get(domain)
	if err != nil {
		return err
	}

	c := config.GetSwiftConnection()

	containerName := swiftV1ContainerPrefixCozy + domain
	containerMeta := &swift.Metadata{"cozy-v1-migrated": "1"}
	err = c.ContainerUpdate(containerName, containerMeta.ContainerHeaders())
	if err != nil {
		return err
	}

	if swiftCluster == 0 {
		swiftCluster = 1
	}
	inst.SwiftCluster = swiftCluster
	return instance.Update(inst)
}

func readObjects(c *swift.Connection, objc chan object,
	containerSrc, containerDst string) error {
	return c.ObjectsWalk(containerSrc, nil, func(opts *swift.ObjectsOpts) (interface{}, error) {
		objs, err := c.Objects(containerSrc, opts)
		if err != nil {
			return nil, err
		}
		for _, obj := range objs {
			objc <- object{
				obj:          obj,
				containerSrc: containerSrc,
				containerDst: containerDst,
			}
		}
		return objs, err
	})
}

func copyObjects(c *swift.Connection, db couchdb.Database,
	objc chan object,
	errc chan error) {

	copyBuffer := make([]byte, 128*1024)

	for obj := range objc {
		var err error
		containerSrc := obj.containerSrc
		containerDst := obj.containerDst
		switch {
		case strings.HasPrefix(containerSrc, swiftV1ContainerPrefixCozy):
			err = copyFileDataObject(c, db, containerSrc, containerDst, obj.obj, copyBuffer)
		case strings.HasPrefix(containerSrc, swiftV1ContainerPrefixData):
			err = copyThumbnailDataObject(c, db, containerSrc, containerDst, obj.obj, copyBuffer)
		}
		if err != nil {
			errc <- err
		}
	}

	errc <- nil
}

func copyFileDataObject(c *swift.Connection, db couchdb.Database,
	containerSrc, containerDst string,
	objSrc swift.Object,
	copyBuffer []byte) error {
	if objSrc.ContentType == dirContentType {
		return nil
	}
	dirID, name, ok := splitV2ObjectName(objSrc.Name)
	if !ok {
		return nil
	}
	var res couchdb.ViewResponse
	err := couchdb.ExecView(db, consts.FilesByParentView, &couchdb.ViewRequest{
		Key:         []string{dirID, consts.FileType, name},
		IncludeDocs: false,
	}, &res)
	if err != nil {
		return err
	}
	if len(res.Rows) == 0 {
		return os.ErrNotExist
	}
	objNameDst := vfsswift.MakeObjectName(res.Rows[0].ID)
	return copyObject(c, db, containerSrc, containerDst, objSrc, objNameDst, copyBuffer)
}

func copyThumbnailDataObject(c *swift.Connection, db couchdb.Database,
	containerSrc, containerDst string,
	objSrc swift.Object,
	copyBuffer []byte) error {

	split := strings.SplitN(strings.TrimPrefix(objSrc.Name, "thumbs/"), "-", 2)
	if len(split) != 2 {
		return nil
	}
	objNameDst := "thumbs/" + vfsswift.MakeObjectName(split[0]) + "-" + split[1]
	return copyObject(c, db, containerSrc, containerDst, objSrc, objNameDst, copyBuffer)
}

func copyObject(c *swift.Connection, db couchdb.Database,
	containerSrc, containerDst string,
	objSrc swift.Object, objNameDst string,
	copyBuffer []byte) (err error) {
	infosDst, _, err := c.Object(containerDst, objNameDst)
	if err == nil && infosDst.Hash == objSrc.Hash {
		return nil
	}
	if err != swift.ObjectNotFound {
		return err
	}

	src, h, err := c.ObjectOpen(containerSrc, objSrc.Name, false, nil)
	if err != nil {
		return err
	}
	defer src.Close()

	dst, err := c.ObjectCreate(containerDst, objNameDst, true, objSrc.Hash, objSrc.ContentType, h)
	if err != nil {
		return err
	}
	defer func() {
		if errc := dst.Close(); errc != nil && err == nil {
			err = errc
		}
	}()

	_, err = io.CopyBuffer(dst, src, copyBuffer)
	return err
}

func splitV2ObjectName(objName string) (dirID string, name string, ok bool) {
	split := strings.SplitN(objName, "/", 2)
	if len(split) != 2 {
		return
	}
	dirID, name = split[0], split[1]
	ok = true
	return
}
<|endoftext|>


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

// gnostic_go_generator is a sample Gnostic plugin that generates Go
// code that supports an API.
package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"

	"github.com/golang/protobuf/proto"

	openapiv2 "github.com/googleapis/gnostic/OpenAPIv2"
	openapiv3 "github.com/googleapis/gnostic/OpenAPIv3"
	plugins "github.com/googleapis/gnostic/plugins"
)

var outputPath string // if nonempty, the plugin is run standalone

// respondAndExitIfError checks an error and if it is non-nil, records it and serializes and returns the response and then exits.
func respondAndExitIfError(err error, response *plugins.Response) {
	if err != nil {
		response.Errors = append(response.Errors, err.Error())
		respondAndExit(response)
	}
}

// respondAndExit serializes and returns the plugin response and then exits.
func respondAndExit(response *plugins.Response) {
	if outputPath != "" {
		err := plugins.HandleResponse(response, outputPath)
		if err != nil {
			log.Printf("%s", err.Error())
		}
	} else {
		responseBytes, _ := proto.Marshal(response)
		os.Stdout.Write(responseBytes)
	}
	os.Exit(0)
}

// This is the main function for the code generation plugin.
func main() {
	invocation := os.Args[0]

	// Use the name used to run the plugin to decide which files to generate.
	var files []string
	switch {
	case strings.Contains(invocation, "gnostic-go-client"):
		files = []string{"client.go", "types.go", "constants.go"}
	case strings.Contains(invocation, "gnostic-go-server"):
		files = []string{"server.go", "provider.go", "types.go", "constants.go"}
	default:
		files = []string{"client.go", "server.go", "provider.go", "types.go", "constants.go"}
	}

	// Initialize the plugin response.
	response := &plugins.Response{}
	var packageName string
	var documentv2 *openapiv2.Document
	var documentv3 *openapiv3.Document

	var version string
	var data []byte
	var err error

	if len(os.Args) == 1 {
		// Read the plugin input.
		data, err := ioutil.ReadAll(os.Stdin)
		respondAndExitIfError(err, response)
		if len(data) == 0 {
			respondAndExitIfError(fmt.Errorf("no input data"), response)
		}

		// Deserialize the input.
		request := &plugins.Request{}
		err = proto.Unmarshal(data, request)
		respondAndExitIfError(err, response)

		// Collect parameters passed to the plugin.
		parameters := request.Parameters
		packageName = request.OutputPath // the default package name is the output directory
		for _, parameter := range parameters {
			invocation += " " + parameter.Name + "=" + parameter.Value
			if parameter.Name == "package" {
				packageName = parameter.Value
			}
		}

		// Log the invocation.
		log.Printf("Running %s(input:%s)", invocation, request.Wrapper.Version)

		// Read the document sent by the plugin.
		version = request.Wrapper.Version
		data = request.Wrapper.Value
	} else {
		input := flag.String("input", "", "OpenAPI input (pb)")
		output := flag.String("output", "-", "output path")
		version2 := flag.Bool("v2", false, "OpenAPI version 2")
		version3 := flag.Bool("v3", false, "OpenAPI version 3")

		flag.Parse()
		outputPath = *output
		packageName = outputPath

		switch {
		case *version2:
			version = "v2"
		case *version3:
			version = "v3"
		default:
			version = "v2"
		}

		// Read the input document.
		data, err = ioutil.ReadFile(*input)
		if len(data) == 0 {
			respondAndExitIfError(fmt.Errorf("no input data"), response)
		}
	}

	switch version {
	case "v2":
		documentv2 = &openapiv2.Document{}
		err = proto.Unmarshal(data, documentv2)
		respondAndExitIfError(err, response)
	case "v3":
		documentv3 = &openapiv3.Document{}
		err = proto.Unmarshal(data, documentv3)
		respondAndExitIfError(err, response)
	default:
		err = fmt.Errorf("Unsupported OpenAPI version %s", version)
		respondAndExitIfError(err, response)
	}

	// Create the model.
	var model *ServiceModel
	if documentv2 != nil {
		model, err = NewServiceModelV2(documentv2, packageName)
	} else {
		model, err = NewServiceModelV3(documentv3, packageName)
	}
	respondAndExitIfError(err, response)

	// Create the renderer.
	renderer, err := NewServiceRenderer(model)
	respondAndExitIfError(err, response)

	// Run the renderer to generate files and add them to the response object.
	err = renderer.Generate(response, files)
	respondAndExitIfError(err, response)

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

// gnostic_go_generator is a sample Gnostic plugin that generates Go
// code that supports an API.
package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"

	"github.com/golang/protobuf/proto"

	openapiv2 "github.com/googleapis/gnostic/OpenAPIv2"
	openapiv3 "github.com/googleapis/gnostic/OpenAPIv3"
	plugins "github.com/googleapis/gnostic/plugins"
)

var outputPath string // if nonempty, the plugin is run standalone

// respondAndExitIfError checks an error and if it is non-nil, records it and serializes and returns the response and then exits.
func respondAndExitIfError(err error, response *plugins.Response) {
	if err != nil {
		response.Errors = append(response.Errors, err.Error())
		respondAndExit(response)
	}
}

// respondAndExit serializes and returns the plugin response and then exits.
func respondAndExit(response *plugins.Response) {
	if outputPath != "" {
		err := plugins.HandleResponse(response, outputPath)
		if err != nil {
			log.Printf("%s", err.Error())
		}
	} else {
		responseBytes, _ := proto.Marshal(response)
		os.Stdout.Write(responseBytes)
	}
	os.Exit(0)
}

// This is the main function for the code generation plugin.
func main() {
	invocation := os.Args[0]

	// Use the name used to run the plugin to decide which files to generate.
	var files []string
	switch {
	case strings.Contains(invocation, "gnostic-go-client"):
		files = []string{"client.go", "types.go", "constants.go"}
	case strings.Contains(invocation, "gnostic-go-server"):
		files = []string{"server.go", "provider.go", "types.go", "constants.go"}
	default:
		files = []string{"client.go", "server.go", "provider.go", "types.go", "constants.go"}
	}

	// Initialize the plugin response.
	response := &plugins.Response{}
	var packageName string
	var documentv2 *openapiv2.Document
	var documentv3 *openapiv3.Document

	var version string
	var apiData []byte
	var err error

	if len(os.Args) == 1 {
		// Read the plugin input.
		pluginData, err := ioutil.ReadAll(os.Stdin)
		respondAndExitIfError(err, response)
		if len(pluginData) == 0 {
			respondAndExitIfError(fmt.Errorf("no input data"), response)
		}

		// Deserialize the input.
		request := &plugins.Request{}
		err = proto.Unmarshal(pluginData, request)
		respondAndExitIfError(err, response)

		// Collect parameters passed to the plugin.
		parameters := request.Parameters
		packageName = request.OutputPath // the default package name is the output directory
		for _, parameter := range parameters {
			invocation += " " + parameter.Name + "=" + parameter.Value
			if parameter.Name == "package" {
				packageName = parameter.Value
			}
		}

		// Log the invocation.
		log.Printf("Running plugin %s(input:%s)", invocation, request.Wrapper.Version)

		// Read the document sent by the plugin.
		version = request.Wrapper.Version
		apiData = request.Wrapper.Value
	} else {
		input := flag.String("input", "", "OpenAPI input (pb)")
		output := flag.String("output", "-", "output path")
		version2 := flag.Bool("v2", false, "OpenAPI version 2")
		version3 := flag.Bool("v3", false, "OpenAPI version 3")

		flag.Parse()
		outputPath = *output
		packageName = outputPath

		switch {
		case *version2:
			version = "v2"
		case *version3:
			version = "v3"
		default:
			version = "v2"
		}

		// Read the input document.
		apiData, err = ioutil.ReadFile(*input)
		if len(apiData) == 0 {
			respondAndExitIfError(fmt.Errorf("no input data"), response)
		}
	}

	switch version {
	case "v2":
		documentv2 = &openapiv2.Document{}
		err = proto.Unmarshal(apiData, documentv2)
		respondAndExitIfError(err, response)
	case "v3":
		documentv3 = &openapiv3.Document{}
		err = proto.Unmarshal(apiData, documentv3)
		respondAndExitIfError(err, response)
	default:
		err = fmt.Errorf("Unsupported OpenAPI version %s", version)
		respondAndExitIfError(err, response)
	}

	// Create the model.
	var model *ServiceModel
	if documentv2 != nil {
		model, err = NewServiceModelV2(documentv2, packageName)
	} else {
		model, err = NewServiceModelV3(documentv3, packageName)
	}
	respondAndExitIfError(err, response)

	// Create the renderer.
	renderer, err := NewServiceRenderer(model)
	respondAndExitIfError(err, response)

	// Run the renderer to generate files and add them to the response object.
	err = renderer.Generate(response, files)
	respondAndExitIfError(err, response)

	// Return with success.
	respondAndExit(response)
}
<|endoftext|>


<commit_before>// Copyright © 2016 Mark Bates <mark@markbates.com>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"errors"
	"fmt"
	"html/template"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/spf13/cobra"
)

var force bool
var skipPop bool
var dbType = "postgres"

var newCmd = &cobra.Command{
	Use:   "new [name]",
	Short: "Creates a new Buffalo application",
	RunE: func(cmd *cobra.Command, args []string) error {
		if len(args) == 0 {
			return errors.New("You must enter a name for your new application.")
		}
		name := args[0]
		pwd, err := os.Getwd()
		if err != nil {
			return err
		}
		rootPath := filepath.Join(pwd, name)

		s, _ := os.Stat(rootPath)
		if s != nil {
			if force {
				os.RemoveAll(rootPath)
			} else {
				return fmt.Errorf("%s already exists! Either delete it or use the -f flag to force.\n", name)
			}
		}

		fmt.Printf("-- ./%s\n", name)
		err = os.MkdirAll(name, 0777)
		if err != nil {
			return err
		}

		err = genNewFiles(name, rootPath)
		if err != nil {
			return err
		}

		err = installDeps(pwd, rootPath)
		if err != nil {
			return err
		}
		return err
	},
}

func installDeps(pwd string, rootPath string) error {
	defer os.Chdir(pwd)
	err := os.Chdir(rootPath)
	if err != nil {
		return err
	}

	cmds := []*exec.Cmd{
		exec.Command("go", "get", "-u", "-v", "github.com/Masterminds/glide"),
		exec.Command("go", "install", "-v", "github.com/markbates/refresh"),
		exec.Command("go", "install", "-v", "github.com/markbates/grift"),
		exec.Command("glide", "init", "--non-interactive"),
		exec.Command("glide", "get", "-v", "-u", "--non-interactive", "github.com/markbates/refresh"),
		exec.Command("glide", "get", "-v", "-u", "--non-interactive", "github.com/markbates/grift"),
		exec.Command("refresh", "init"),
	}

	if !skipPop {
		cmds = append(cmds,
			exec.Command("glide", "get", "-v", "-u", "--non-interactive", "github.com/markbates/pop/"),
			exec.Command("glide", "get", "-v", "-u", "--non-interactive", "github.com/markbates/pop/soda"),
			exec.Command("go", "install", "-v", "github.com/markbates/pop/soda"),
			exec.Command("soda", "g", "config", "-t", dbType),
		)
	}

	err = runCommands(cmds...)

	if err != nil {
		return err
	}

	return err
}

func runCommands(cmds ...*exec.Cmd) error {
	for _, cmd := range cmds {
		cmd.Stdin = os.Stdin
		cmd.Stderr = os.Stderr
		cmd.Stdout = os.Stdout
		err := cmd.Run()
		if err != nil {
			return err
		}
	}
	return nil
}

func genNewFiles(name, rootPath string) error {
	packagePath := strings.Replace(rootPath, filepath.Join(os.Getenv("GOPATH"), "src")+"/", "", 1)

	data := map[string]interface{}{
		"name":        name,
		"packagePath": packagePath,
		"actionsPath": filepath.Join(packagePath, "actions"),
	}

	for fn, tv := range newTemplates {
		dir := filepath.Dir(fn)
		err := os.MkdirAll(filepath.Join(rootPath, dir), 0777)
		if err != nil {
			return err
		}
		t, err := template.New(fn).Parse(tv)
		if err != nil {
			return err
		}
		fmt.Printf("-- ./%s/%s\n", name, fn)
		f, err := os.Create(filepath.Join(rootPath, fn))
		if err != nil {
			return err
		}
		err = t.Execute(f, data)
		if err != nil {
			return err
		}
	}
	return nil
}

func init() {
	RootCmd.AddCommand(newCmd)

	// Here you will define your flags and configuration settings.

	// Cobra supports Persistent Flags which will work for this command
	// and all subcommands, e.g.:
	// newCmd.PersistentFlags().String("foo", "", "A help for foo")

	// Cobra supports local flags which will only run when this command
	// is called directly, e.g.:
	newCmd.Flags().BoolVarP(&force, "force", "f", false, "delete and remake if the app already exists")
	newCmd.Flags().BoolVar(&skipPop, "skip-pop", false, "skips add pop/soda to your app")
	newCmd.Flags().StringVar(&dbType, "db-type", "postgres", "specify the type of database you want to use [postgres, mysql, sqlite3]")

}
<commit_msg>trying to get this dep installer right<commit_after>// Copyright © 2016 Mark Bates <mark@markbates.com>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"errors"
	"fmt"
	"html/template"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/spf13/cobra"
)

var force bool
var skipPop bool
var dbType = "postgres"

var newCmd = &cobra.Command{
	Use:   "new [name]",
	Short: "Creates a new Buffalo application",
	RunE: func(cmd *cobra.Command, args []string) error {
		if len(args) == 0 {
			return errors.New("You must enter a name for your new application.")
		}
		name := args[0]
		pwd, err := os.Getwd()
		if err != nil {
			return err
		}
		rootPath := filepath.Join(pwd, name)

		s, _ := os.Stat(rootPath)
		if s != nil {
			if force {
				os.RemoveAll(rootPath)
			} else {
				return fmt.Errorf("%s already exists! Either delete it or use the -f flag to force.\n", name)
			}
		}

		fmt.Printf("-- ./%s\n", name)
		err = os.MkdirAll(name, 0777)
		if err != nil {
			return err
		}

		err = genNewFiles(name, rootPath)
		if err != nil {
			return err
		}

		err = installDeps(pwd, rootPath)
		if err != nil {
			return err
		}
		return err
	},
}

func installDeps(pwd string, rootPath string) error {
	defer os.Chdir(pwd)
	err := os.Chdir(rootPath)
	if err != nil {
		return err
	}

	cmds := []*exec.Cmd{
		exec.Command("go", "get", "-u", "-v", "github.com/Masterminds/glide"),
		exec.Command("go", "get", "-u", "-v", "github.com/markbates/refresh"),
		exec.Command("go", "get", "-u", "-v", "github.com/markbates/grift"),
		exec.Command("glide", "init", "--non-interactive"),
		exec.Command("glide", "get", "-v", "-u", "--non-interactive", "github.com/markbates/refresh"),
		exec.Command("glide", "get", "-v", "-u", "--non-interactive", "github.com/markbates/grift"),
		exec.Command("refresh", "init"),
	}

	if !skipPop {
		cmds = append(cmds,
			exec.Command("go", "get", "-u", "-v", "github.com/markbates/pop/soda"),
			exec.Command("glide", "get", "-v", "-u", "--non-interactive", "github.com/markbates/pop/"),
			exec.Command("glide", "get", "-v", "-u", "--non-interactive", "github.com/markbates/pop/soda"),
			exec.Command("go", "install", "-v", "github.com/markbates/pop/soda"),
			exec.Command("soda", "g", "config", "-t", dbType),
		)
	}

	err = runCommands(cmds...)

	if err != nil {
		return err
	}

	return err
}

func runCommands(cmds ...*exec.Cmd) error {
	for _, cmd := range cmds {
		cmd.Stdin = os.Stdin
		cmd.Stderr = os.Stderr
		cmd.Stdout = os.Stdout
		err := cmd.Run()
		if err != nil {
			return err
		}
	}
	return nil
}

func genNewFiles(name, rootPath string) error {
	packagePath := strings.Replace(rootPath, filepath.Join(os.Getenv("GOPATH"), "src")+"/", "", 1)

	data := map[string]interface{}{
		"name":        name,
		"packagePath": packagePath,
		"actionsPath": filepath.Join(packagePath, "actions"),
	}

	for fn, tv := range newTemplates {
		dir := filepath.Dir(fn)
		err := os.MkdirAll(filepath.Join(rootPath, dir), 0777)
		if err != nil {
			return err
		}
		t, err := template.New(fn).Parse(tv)
		if err != nil {
			return err
		}
		fmt.Printf("-- ./%s/%s\n", name, fn)
		f, err := os.Create(filepath.Join(rootPath, fn))
		if err != nil {
			return err
		}
		err = t.Execute(f, data)
		if err != nil {
			return err
		}
	}
	return nil
}

func init() {
	RootCmd.AddCommand(newCmd)

	// Here you will define your flags and configuration settings.

	// Cobra supports Persistent Flags which will work for this command
	// and all subcommands, e.g.:
	// newCmd.PersistentFlags().String("foo", "", "A help for foo")

	// Cobra supports local flags which will only run when this command
	// is called directly, e.g.:
	newCmd.Flags().BoolVarP(&force, "force", "f", false, "delete and remake if the app already exists")
	newCmd.Flags().BoolVar(&skipPop, "skip-pop", false, "skips add pop/soda to your app")
	newCmd.Flags().StringVar(&dbType, "db-type", "postgres", "specify the type of database you want to use [postgres, mysql, sqlite3]")

}
<|endoftext|>


<commit_before>package peer

import (
	"net"
	"time"

	"github.com/cenkalti/rain/torrent/internal/peerconn"
	"github.com/cenkalti/rain/torrent/internal/peerprotocol"
	"github.com/cenkalti/rain/torrent/internal/pexlist"
)

type pex struct {
	conn  *peerconn.Conn
	extID uint8

	// Contains added and dropped peers.
	pexList *pexlist.PEXList

	// To send connected peers at interval
	pexTicker *time.Ticker

	pexAddPeerC  chan *net.TCPAddr
	pexDropPeerC chan *net.TCPAddr

	closeC chan struct{}
	doneC  chan struct{}
}

func newPEX(conn *peerconn.Conn, extID uint8, initialPeers map[*Peer]struct{}) *pex {
	pl := pexlist.New()
	for pe := range initialPeers {
		if pe.Addr().String() != conn.Addr().String() {
			pl.Add(pe.Addr())
		}
	}
	return &pex{
		conn:         conn,
		extID:        extID,
		pexList:      pl,
		pexAddPeerC:  make(chan *net.TCPAddr),
		pexDropPeerC: make(chan *net.TCPAddr),
		closeC:       make(chan struct{}),
		doneC:        make(chan struct{}),
	}
}

func (p *pex) close() {
	close(p.closeC)
	<-p.doneC
}

func (p *pex) run() {
	defer close(p.doneC)

	p.pexFlushPeers()

	p.pexTicker = time.NewTicker(time.Minute)
	defer p.pexTicker.Stop()

	for {
		select {
		case addr := <-p.pexAddPeerC:
			p.pexList.Add(addr)
		case addr := <-p.pexDropPeerC:
			p.pexList.Drop(addr)
		case <-p.pexTicker.C:
			p.pexFlushPeers()
		case <-p.closeC:
			return
		}
	}
}

func (p *pex) Add(addr *net.TCPAddr) {
	select {
	case p.pexAddPeerC <- addr:
	case <-p.doneC:
	}
}

func (p *pex) Drop(addr *net.TCPAddr) {
	select {
	case p.pexDropPeerC <- addr:
	case <-p.doneC:
	}
}

func (p *pex) pexFlushPeers() {
	added, dropped := p.pexList.Flush()
	extPEXMsg := peerprotocol.ExtensionPEXMessage{
		Added:   added,
		Dropped: dropped,
	}
	msg := peerprotocol.ExtensionMessage{
		ExtendedMessageID: p.extID,
		Payload:           extPEXMsg,
	}
	p.conn.SendMessage(msg)
}
<commit_msg>do not send empty pex message<commit_after>package peer

import (
	"net"
	"time"

	"github.com/cenkalti/rain/torrent/internal/peerconn"
	"github.com/cenkalti/rain/torrent/internal/peerprotocol"
	"github.com/cenkalti/rain/torrent/internal/pexlist"
)

type pex struct {
	conn  *peerconn.Conn
	extID uint8

	// Contains added and dropped peers.
	pexList *pexlist.PEXList

	// To send connected peers at interval
	pexTicker *time.Ticker

	pexAddPeerC  chan *net.TCPAddr
	pexDropPeerC chan *net.TCPAddr

	closeC chan struct{}
	doneC  chan struct{}
}

func newPEX(conn *peerconn.Conn, extID uint8, initialPeers map[*Peer]struct{}) *pex {
	pl := pexlist.New()
	for pe := range initialPeers {
		if pe.Addr().String() != conn.Addr().String() {
			pl.Add(pe.Addr())
		}
	}
	return &pex{
		conn:         conn,
		extID:        extID,
		pexList:      pl,
		pexAddPeerC:  make(chan *net.TCPAddr),
		pexDropPeerC: make(chan *net.TCPAddr),
		closeC:       make(chan struct{}),
		doneC:        make(chan struct{}),
	}
}

func (p *pex) close() {
	close(p.closeC)
	<-p.doneC
}

func (p *pex) run() {
	defer close(p.doneC)

	p.pexFlushPeers()

	p.pexTicker = time.NewTicker(time.Minute)
	defer p.pexTicker.Stop()

	for {
		select {
		case addr := <-p.pexAddPeerC:
			p.pexList.Add(addr)
		case addr := <-p.pexDropPeerC:
			p.pexList.Drop(addr)
		case <-p.pexTicker.C:
			p.pexFlushPeers()
		case <-p.closeC:
			return
		}
	}
}

func (p *pex) Add(addr *net.TCPAddr) {
	select {
	case p.pexAddPeerC <- addr:
	case <-p.doneC:
	}
}

func (p *pex) Drop(addr *net.TCPAddr) {
	select {
	case p.pexDropPeerC <- addr:
	case <-p.doneC:
	}
}

func (p *pex) pexFlushPeers() {
	added, dropped := p.pexList.Flush()
	if len(added) == 0 && len(dropped) == 0 {
		return
	}
	extPEXMsg := peerprotocol.ExtensionPEXMessage{
		Added:   added,
		Dropped: dropped,
	}
	msg := peerprotocol.ExtensionMessage{
		ExtendedMessageID: p.extID,
		Payload:           extPEXMsg,
	}
	p.conn.SendMessage(msg)
}
<|endoftext|>