<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 components

import (
	"strings"

	"k8s.io/klog"
	"k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/loader"
)

// KubeletOptionsBuilder adds options for kubelets
type KubeletOptionsBuilder struct {
	Context *OptionsContext
}

var _ loader.OptionsBuilder = &KubeletOptionsBuilder{}

// BuildOptions is responsible for filling the defaults for the kubelet
func (b *KubeletOptionsBuilder) BuildOptions(o interface{}) error {
	clusterSpec := o.(*kops.ClusterSpec)

	kubernetesVersion, err := KubernetesVersion(clusterSpec)
	if err != nil {
		return err
	}

	if clusterSpec.Kubelet == nil {
		clusterSpec.Kubelet = &kops.KubeletConfigSpec{}
	}
	if clusterSpec.MasterKubelet == nil {
		clusterSpec.MasterKubelet = &kops.KubeletConfigSpec{}
	}

	if clusterSpec.KubeAPIServer != nil && clusterSpec.KubeAPIServer.EnableBootstrapAuthToken != nil {
		if *clusterSpec.KubeAPIServer.EnableBootstrapAuthToken {
			if clusterSpec.Kubelet.BootstrapKubeconfig == "" {
				clusterSpec.Kubelet.BootstrapKubeconfig = "/var/lib/kubelet/bootstrap-kubeconfig"
			}
		}
	}

	// Standard options
	clusterSpec.Kubelet.EnableDebuggingHandlers = fi.Bool(true)
	clusterSpec.Kubelet.PodManifestPath = "/etc/kubernetes/manifests"
	clusterSpec.Kubelet.LogLevel = fi.Int32(2)
	clusterSpec.Kubelet.ClusterDomain = clusterSpec.ClusterDNSDomain
	clusterSpec.Kubelet.NonMasqueradeCIDR = clusterSpec.NonMasqueradeCIDR

	// AllowPrivileged is deprecated and removed in v1.14.
	// See https://github.com/kubernetes/kubernetes/pull/71835
	if kubernetesVersion.Major == 1 && kubernetesVersion.Minor >= 14 {
		if clusterSpec.Kubelet.AllowPrivileged != nil {
			// If it is explicitly set to false, return an error, because this
			// behavior is no longer supported in v1.14 (the default was true, prior).
			if *clusterSpec.Kubelet.AllowPrivileged == false {
				klog.Warningf("Kubelet's --allow-privileged flag is no longer supported in v1.14.")
			}
			// Explicitly set it to nil, so it won't be passed on the command line.
			clusterSpec.Kubelet.AllowPrivileged = nil
		}
	} else {
		clusterSpec.Kubelet.AllowPrivileged = fi.Bool(true)
	}

	if clusterSpec.Kubelet.ClusterDNS == "" {
		ip, err := WellKnownServiceIP(clusterSpec, 10)
		if err != nil {
			return err
		}
		clusterSpec.Kubelet.ClusterDNS = ip.String()
	}

	if b.Context.IsKubernetesLT("1.7") {
		// babysit-daemons removed in 1.7
		clusterSpec.Kubelet.BabysitDaemons = fi.Bool(true)
	}

	clusterSpec.MasterKubelet.RegisterSchedulable = fi.Bool(false)
	// Replace the CIDR with a CIDR allocated by KCM (the default, but included for clarity)
	// We _do_ allow debugging handlers, so we can do logs
	// This does allow more access than we would like though
	clusterSpec.MasterKubelet.EnableDebuggingHandlers = fi.Bool(true)

	// In 1.5 we fixed this, but in 1.4 we need to set the PodCIDR on the master
	// so that hostNetwork pods can come up
	if kubernetesVersion.Major == 1 && kubernetesVersion.Minor <= 4 {
		// We bootstrap with a fake CIDR, but then this will be replaced (unless we're running with _isolated_master)
		clusterSpec.MasterKubelet.PodCIDR = "10.123.45.0/28"
	}

	// 1.5 deprecates the reconcile cidr option (and 1.6 removes it)
	if kubernetesVersion.Major == 1 && kubernetesVersion.Minor <= 4 {
		clusterSpec.MasterKubelet.ReconcileCIDR = fi.Bool(true)

		if fi.BoolValue(clusterSpec.IsolateMasters) {
			clusterSpec.MasterKubelet.ReconcileCIDR = fi.Bool(false)
		}

		usesKubenet, err := UsesKubenet(clusterSpec)
		if err != nil {
			return err
		}
		if usesKubenet {
			clusterSpec.Kubelet.ReconcileCIDR = fi.Bool(true)
		}
	}

	if kubernetesVersion.Major == 1 && kubernetesVersion.Minor >= 4 {
		// For pod eviction in low memory or empty disk situations
		if clusterSpec.Kubelet.EvictionHard == nil {
			evictionHard := []string{
				// TODO: Some people recommend 250Mi, but this would hurt small machines
				"memory.available<100Mi",

				// Disk eviction (evict old images)
				// We don't need to specify both, but it seems harmless / safer
				"nodefs.available<10%",
				"nodefs.inodesFree<5%",
				"imagefs.available<10%",
				"imagefs.inodesFree<5%",
			}
			clusterSpec.Kubelet.EvictionHard = fi.String(strings.Join(evictionHard, ","))
		}
	}

	if b.Context.IsKubernetesGTE("1.6") {
		// for 1.6+ use kubeconfig instead of api-servers
		const kubeconfigPath = "/var/lib/kubelet/kubeconfig"
		clusterSpec.Kubelet.KubeconfigPath = kubeconfigPath
		clusterSpec.MasterKubelet.KubeconfigPath = kubeconfigPath

		// Only pass require-kubeconfig to versions prior to 1.9; deprecated & being removed
		if b.Context.IsKubernetesLT("1.9") {
			clusterSpec.Kubelet.RequireKubeconfig = fi.Bool(true)
			clusterSpec.MasterKubelet.RequireKubeconfig = fi.Bool(true)
		}
	} else {
		// Legacy behaviour for <= 1.5
		clusterSpec.Kubelet.APIServers = "https://" + clusterSpec.MasterInternalName
		clusterSpec.MasterKubelet.APIServers = "http://127.0.0.1:8080"
	}

	// IsolateMasters enables the legacy behaviour, where master pods on a separate network
	// In newer versions of kubernetes, most of that functionality has been removed though
	if fi.BoolValue(clusterSpec.IsolateMasters) {
		clusterSpec.MasterKubelet.EnableDebuggingHandlers = fi.Bool(false)
		clusterSpec.MasterKubelet.HairpinMode = "none"
	}

	cloudProvider := kops.CloudProviderID(clusterSpec.CloudProvider)

	clusterSpec.Kubelet.CgroupRoot = "/"

	klog.V(1).Infof("Cloud Provider: %s", cloudProvider)
	if cloudProvider == kops.CloudProviderAWS {
		clusterSpec.Kubelet.CloudProvider = "aws"

		// For 1.6 we're using much cleaner cgroup hierarchies
		// but we keep the settings we've tested for k8s 1.5 and lower
		// (see https://github.com/kubernetes/kubernetes/pull/41349)
		if kubernetesVersion.Major == 1 && kubernetesVersion.Minor <= 5 {
			clusterSpec.Kubelet.CgroupRoot = "docker"
		}

		// Use the hostname from the AWS metadata service
		// if hostnameOverride is not set.
		if clusterSpec.Kubelet.HostnameOverride == "" {
			clusterSpec.Kubelet.HostnameOverride = "@aws"
		}
	}

	if cloudProvider == kops.CloudProviderDO {
		clusterSpec.Kubelet.CloudProvider = "external"
		clusterSpec.Kubelet.HostnameOverride = "@digitalocean"
	}

	if cloudProvider == kops.CloudProviderGCE {
		clusterSpec.Kubelet.CloudProvider = "gce"
		clusterSpec.Kubelet.HairpinMode = "promiscuous-bridge"

		if clusterSpec.CloudConfig == nil {
			clusterSpec.CloudConfig = &kops.CloudConfiguration{}
		}
		clusterSpec.CloudConfig.Multizone = fi.Bool(true)
		clusterSpec.CloudConfig.NodeTags = fi.String(GCETagForRole(b.Context.ClusterName, kops.InstanceGroupRoleNode))
	}

	if cloudProvider == kops.CloudProviderVSphere {
		clusterSpec.Kubelet.CloudProvider = "vsphere"
		clusterSpec.Kubelet.HairpinMode = "promiscuous-bridge"
	}

	if cloudProvider == kops.CloudProviderOpenstack {
		clusterSpec.Kubelet.CloudProvider = "openstack"
	}

	if cloudProvider == kops.CloudProviderALI {
		clusterSpec.Kubelet.CloudProvider = "alicloud"
	}

	if clusterSpec.ExternalCloudControllerManager != nil {
		clusterSpec.Kubelet.CloudProvider = "external"
	}

	usesKubenet, err := UsesKubenet(clusterSpec)
	if err != nil {
		return err
	}
	if usesKubenet {
		clusterSpec.Kubelet.NetworkPluginName = "kubenet"

		if kubernetesVersion.Major == 1 && kubernetesVersion.Minor >= 4 {
			// AWS MTU is 9001
			clusterSpec.Kubelet.NetworkPluginMTU = fi.Int32(9001)
		}
	}

	// Specify our pause image
	image := "k8s.gcr.io/pause-amd64:3.0"
	if image, err = b.Context.AssetBuilder.RemapImage(image); err != nil {
		return err
	}
	clusterSpec.Kubelet.PodInfraContainerImage = image

	if clusterSpec.Kubelet.FeatureGates == nil {
		clusterSpec.Kubelet.FeatureGates = make(map[string]string)
	}
	if _, found := clusterSpec.Kubelet.FeatureGates["ExperimentalCriticalPodAnnotation"]; !found {
		if b.Context.IsKubernetesGTE("1.5.2") {
			clusterSpec.Kubelet.FeatureGates["ExperimentalCriticalPodAnnotation"] = "true"
		}
	}

	return nil
}
<commit_msg>Dont set CriticalPodAnnotation feature gate in k8s 1.16<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 components

import (
	"strings"

	"k8s.io/klog"
	"k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/loader"
)

// KubeletOptionsBuilder adds options for kubelets
type KubeletOptionsBuilder struct {
	Context *OptionsContext
}

var _ loader.OptionsBuilder = &KubeletOptionsBuilder{}

// BuildOptions is responsible for filling the defaults for the kubelet
func (b *KubeletOptionsBuilder) BuildOptions(o interface{}) error {
	clusterSpec := o.(*kops.ClusterSpec)

	kubernetesVersion, err := KubernetesVersion(clusterSpec)
	if err != nil {
		return err
	}

	if clusterSpec.Kubelet == nil {
		clusterSpec.Kubelet = &kops.KubeletConfigSpec{}
	}
	if clusterSpec.MasterKubelet == nil {
		clusterSpec.MasterKubelet = &kops.KubeletConfigSpec{}
	}

	if clusterSpec.KubeAPIServer != nil && clusterSpec.KubeAPIServer.EnableBootstrapAuthToken != nil {
		if *clusterSpec.KubeAPIServer.EnableBootstrapAuthToken {
			if clusterSpec.Kubelet.BootstrapKubeconfig == "" {
				clusterSpec.Kubelet.BootstrapKubeconfig = "/var/lib/kubelet/bootstrap-kubeconfig"
			}
		}
	}

	// Standard options
	clusterSpec.Kubelet.EnableDebuggingHandlers = fi.Bool(true)
	clusterSpec.Kubelet.PodManifestPath = "/etc/kubernetes/manifests"
	clusterSpec.Kubelet.LogLevel = fi.Int32(2)
	clusterSpec.Kubelet.ClusterDomain = clusterSpec.ClusterDNSDomain
	clusterSpec.Kubelet.NonMasqueradeCIDR = clusterSpec.NonMasqueradeCIDR

	// AllowPrivileged is deprecated and removed in v1.14.
	// See https://github.com/kubernetes/kubernetes/pull/71835
	if kubernetesVersion.Major == 1 && kubernetesVersion.Minor >= 14 {
		if clusterSpec.Kubelet.AllowPrivileged != nil {
			// If it is explicitly set to false, return an error, because this
			// behavior is no longer supported in v1.14 (the default was true, prior).
			if *clusterSpec.Kubelet.AllowPrivileged == false {
				klog.Warningf("Kubelet's --allow-privileged flag is no longer supported in v1.14.")
			}
			// Explicitly set it to nil, so it won't be passed on the command line.
			clusterSpec.Kubelet.AllowPrivileged = nil
		}
	} else {
		clusterSpec.Kubelet.AllowPrivileged = fi.Bool(true)
	}

	if clusterSpec.Kubelet.ClusterDNS == "" {
		ip, err := WellKnownServiceIP(clusterSpec, 10)
		if err != nil {
			return err
		}
		clusterSpec.Kubelet.ClusterDNS = ip.String()
	}

	if b.Context.IsKubernetesLT("1.7") {
		// babysit-daemons removed in 1.7
		clusterSpec.Kubelet.BabysitDaemons = fi.Bool(true)
	}

	clusterSpec.MasterKubelet.RegisterSchedulable = fi.Bool(false)
	// Replace the CIDR with a CIDR allocated by KCM (the default, but included for clarity)
	// We _do_ allow debugging handlers, so we can do logs
	// This does allow more access than we would like though
	clusterSpec.MasterKubelet.EnableDebuggingHandlers = fi.Bool(true)

	// In 1.5 we fixed this, but in 1.4 we need to set the PodCIDR on the master
	// so that hostNetwork pods can come up
	if kubernetesVersion.Major == 1 && kubernetesVersion.Minor <= 4 {
		// We bootstrap with a fake CIDR, but then this will be replaced (unless we're running with _isolated_master)
		clusterSpec.MasterKubelet.PodCIDR = "10.123.45.0/28"
	}

	// 1.5 deprecates the reconcile cidr option (and 1.6 removes it)
	if kubernetesVersion.Major == 1 && kubernetesVersion.Minor <= 4 {
		clusterSpec.MasterKubelet.ReconcileCIDR = fi.Bool(true)

		if fi.BoolValue(clusterSpec.IsolateMasters) {
			clusterSpec.MasterKubelet.ReconcileCIDR = fi.Bool(false)
		}

		usesKubenet, err := UsesKubenet(clusterSpec)
		if err != nil {
			return err
		}
		if usesKubenet {
			clusterSpec.Kubelet.ReconcileCIDR = fi.Bool(true)
		}
	}

	if kubernetesVersion.Major == 1 && kubernetesVersion.Minor >= 4 {
		// For pod eviction in low memory or empty disk situations
		if clusterSpec.Kubelet.EvictionHard == nil {
			evictionHard := []string{
				// TODO: Some people recommend 250Mi, but this would hurt small machines
				"memory.available<100Mi",

				// Disk eviction (evict old images)
				// We don't need to specify both, but it seems harmless / safer
				"nodefs.available<10%",
				"nodefs.inodesFree<5%",
				"imagefs.available<10%",
				"imagefs.inodesFree<5%",
			}
			clusterSpec.Kubelet.EvictionHard = fi.String(strings.Join(evictionHard, ","))
		}
	}

	if b.Context.IsKubernetesGTE("1.6") {
		// for 1.6+ use kubeconfig instead of api-servers
		const kubeconfigPath = "/var/lib/kubelet/kubeconfig"
		clusterSpec.Kubelet.KubeconfigPath = kubeconfigPath
		clusterSpec.MasterKubelet.KubeconfigPath = kubeconfigPath

		// Only pass require-kubeconfig to versions prior to 1.9; deprecated & being removed
		if b.Context.IsKubernetesLT("1.9") {
			clusterSpec.Kubelet.RequireKubeconfig = fi.Bool(true)
			clusterSpec.MasterKubelet.RequireKubeconfig = fi.Bool(true)
		}
	} else {
		// Legacy behaviour for <= 1.5
		clusterSpec.Kubelet.APIServers = "https://" + clusterSpec.MasterInternalName
		clusterSpec.MasterKubelet.APIServers = "http://127.0.0.1:8080"
	}

	// IsolateMasters enables the legacy behaviour, where master pods on a separate network
	// In newer versions of kubernetes, most of that functionality has been removed though
	if fi.BoolValue(clusterSpec.IsolateMasters) {
		clusterSpec.MasterKubelet.EnableDebuggingHandlers = fi.Bool(false)
		clusterSpec.MasterKubelet.HairpinMode = "none"
	}

	cloudProvider := kops.CloudProviderID(clusterSpec.CloudProvider)

	clusterSpec.Kubelet.CgroupRoot = "/"

	klog.V(1).Infof("Cloud Provider: %s", cloudProvider)
	if cloudProvider == kops.CloudProviderAWS {
		clusterSpec.Kubelet.CloudProvider = "aws"

		// For 1.6 we're using much cleaner cgroup hierarchies
		// but we keep the settings we've tested for k8s 1.5 and lower
		// (see https://github.com/kubernetes/kubernetes/pull/41349)
		if kubernetesVersion.Major == 1 && kubernetesVersion.Minor <= 5 {
			clusterSpec.Kubelet.CgroupRoot = "docker"
		}

		// Use the hostname from the AWS metadata service
		// if hostnameOverride is not set.
		if clusterSpec.Kubelet.HostnameOverride == "" {
			clusterSpec.Kubelet.HostnameOverride = "@aws"
		}
	}

	if cloudProvider == kops.CloudProviderDO {
		clusterSpec.Kubelet.CloudProvider = "external"
		clusterSpec.Kubelet.HostnameOverride = "@digitalocean"
	}

	if cloudProvider == kops.CloudProviderGCE {
		clusterSpec.Kubelet.CloudProvider = "gce"
		clusterSpec.Kubelet.HairpinMode = "promiscuous-bridge"

		if clusterSpec.CloudConfig == nil {
			clusterSpec.CloudConfig = &kops.CloudConfiguration{}
		}
		clusterSpec.CloudConfig.Multizone = fi.Bool(true)
		clusterSpec.CloudConfig.NodeTags = fi.String(GCETagForRole(b.Context.ClusterName, kops.InstanceGroupRoleNode))
	}

	if cloudProvider == kops.CloudProviderVSphere {
		clusterSpec.Kubelet.CloudProvider = "vsphere"
		clusterSpec.Kubelet.HairpinMode = "promiscuous-bridge"
	}

	if cloudProvider == kops.CloudProviderOpenstack {
		clusterSpec.Kubelet.CloudProvider = "openstack"
	}

	if cloudProvider == kops.CloudProviderALI {
		clusterSpec.Kubelet.CloudProvider = "alicloud"
	}

	if clusterSpec.ExternalCloudControllerManager != nil {
		clusterSpec.Kubelet.CloudProvider = "external"
	}

	usesKubenet, err := UsesKubenet(clusterSpec)
	if err != nil {
		return err
	}
	if usesKubenet {
		clusterSpec.Kubelet.NetworkPluginName = "kubenet"

		if kubernetesVersion.Major == 1 && kubernetesVersion.Minor >= 4 {
			// AWS MTU is 9001
			clusterSpec.Kubelet.NetworkPluginMTU = fi.Int32(9001)
		}
	}

	// Specify our pause image
	image := "k8s.gcr.io/pause-amd64:3.0"
	if image, err = b.Context.AssetBuilder.RemapImage(image); err != nil {
		return err
	}
	clusterSpec.Kubelet.PodInfraContainerImage = image

	if clusterSpec.Kubelet.FeatureGates == nil {
		clusterSpec.Kubelet.FeatureGates = make(map[string]string)
	}
	if _, found := clusterSpec.Kubelet.FeatureGates["ExperimentalCriticalPodAnnotation"]; !found {
		if b.Context.IsKubernetesGTE("1.5.2") && b.Context.IsKubernetesLT("1.16") {
			clusterSpec.Kubelet.FeatureGates["ExperimentalCriticalPodAnnotation"] = "true"
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package vsphere

import (
	"fmt"
	"io"
	"io/ioutil"
	"net/url"
	"path"
	"runtime"
	"strings"
	"sync"
	"time"

	"gopkg.in/gcfg.v1"

	log "github.com/Sirupsen/logrus"
	"github.com/matt-deboer/etcdcd/pkg/platform"
	"github.com/vmware/govmomi"
	"github.com/vmware/govmomi/find"
	"github.com/vmware/govmomi/object"
	"github.com/vmware/govmomi/property"
	"github.com/vmware/govmomi/session"
	"github.com/vmware/govmomi/vim25"
	"github.com/vmware/govmomi/vim25/mo"
	"github.com/vmware/govmomi/vim25/types"
	"golang.org/x/net/context"

	etcd "github.com/coreos/etcd/client"
)

const (
	ActivePowerState         = "poweredOn"
	RoundTripperDefaultCount = 3
)

var clientLock sync.Mutex

// VSphere is an implementation of cloud provider Interface for VSphere.
type VSphere struct {
	client *govmomi.Client
	cfg    *VSphereConfig
	// InstanceID of the server where this VSphere object is instantiated.
	localInstanceID string
}

type VSphereConfig struct {
	Global struct {
		// vCenter username.
		User string `gcfg:"user"`
		// vCenter password in clear text.
		Password string `gcfg:"password"`
		// vCenter IP.
		VCenterIP string `gcfg:"server"`
		// vCenter port.
		VCenterPort string `gcfg:"port"`
		// True if vCenter uses self-signed cert.
		InsecureFlag bool `gcfg:"insecure-flag"`
		// Datacenter in which VMs are located.
		Datacenter string `gcfg:"datacenter"`
		// Datastore in which vmdks are stored.
		Datastore string `gcfg:"datastore"`
		// WorkingDir is path where VMs can be found.
		WorkingDir string `gcfg:"working-dir"`
		// Soap round tripper count (retries = RoundTripper - 1)
		RoundTripperCount uint `gcfg:"soap-roundtrip-count"`
		// VMUUID is the VM Instance UUID of virtual machine which can be retrieved from instanceUuid
		// property in VmConfigInfo, or also set as vc.uuid in VMX file.
		// If not set, will be fetched from the machine via sysfs (requires root)
		VMUUID string `gcfg:"vm-uuid"`
	}
	Network struct {
		// PublicNetwork is name of the network the VMs are joined to.
		PublicNetwork string `gcfg:"public-network"`
	}
}

func init() {
	platform.Register("vsphere", func(config io.Reader) (platform.Platform, error) {
		cfg, err := readConfig(config)
		if err != nil && !strings.Contains(err.Error(), "warnings") {
			log.Fatal("Failed reading config: ", err)
		}
		return newVSphere(cfg)
	})
}

// Parses vSphere cloud config file and stores it into VSphereConfig.
func readConfig(config io.Reader) (VSphereConfig, error) {
	if config == nil {
		err := fmt.Errorf("no config file given")
		return VSphereConfig{}, err
	}

	var cfg VSphereConfig
	err := gcfg.ReadInto(&cfg, config)
	return cfg, err
}

// ExpectedMembers returns a list of members that should form the cluster
func (vs *VSphere) ExpectedMembers(
	memberFilter string, clientScheme string, clientPort int, serverScheme string, serverPort int) ([]etcd.Member, error) {

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	err := vSphereLogin(ctx, vs)
	if err != nil {
		log.Errorf("VCenter login failed; %v", err)
		return nil, err
	}

	members := []etcd.Member{}
	names, err := vs.list(memberFilter)
	if err != nil {
		return nil, err
	}
	for _, name := range names {
		member := etcd.Member{Name: name, ClientURLs: []string{}, PeerURLs: []string{}}
		for tries := 0; tries <= 10 && len(member.PeerURLs) == 0; tries++ {
			addrs, err := vs.getAddresses(name)
			if err != nil {
				return nil, err
			}
			for _, a := range addrs {
				addr := a
				if strings.Contains(a, ":") {
					addr = "[" + a + "]"
				}
				member.ClientURLs = append(member.ClientURLs, fmt.Sprintf("%s://%s:%d", clientScheme, addr, clientPort))
				member.PeerURLs = append(member.PeerURLs, fmt.Sprintf("%s://%s:%d", serverScheme, addr, serverPort))
			}

			if len(member.PeerURLs) > 0 {
				if log.GetLevel() >= log.DebugLevel {
					log.Debugf("ExpectedMembers: member: %#v", member)
				}
			} else {
				sleepTime := (2 * time.Second)
				if log.GetLevel() >= log.DebugLevel {
					log.Debugf("%s has no addresses yet; sleeping for %s", name, sleepTime)
				}
				time.Sleep(sleepTime)
			}
		}
		members = append(members, member)
	}
	return members, nil
}

// LocalInstanceName returns a list of members that should form the cluster
func (vs *VSphere) LocalInstanceName() string {
	return vs.localInstanceID
}

// Returns the name of the VM on which this code is running.
// Prerequisite: this code assumes VMWare vmtools or open-vm-tools to be installed in the VM.
// Will attempt to determine the machine's name via it's UUID in this precedence order, failing if neither have a UUID:
// * cloud config value VMUUID
// * sysfs entry
func getVMName(client *govmomi.Client, cfg *VSphereConfig) (string, error) {

	var vmUUID string

	if cfg.Global.VMUUID != "" {
		vmUUID = cfg.Global.VMUUID
	} else {
		// This needs root privileges on the host, and will fail otherwise.
		vmUUIDbytes, err := ioutil.ReadFile("/sys/devices/virtual/dmi/id/product_uuid")
		if err != nil {
			return "", err
		}

		vmUUID = string(vmUUIDbytes)
		cfg.Global.VMUUID = vmUUID
	}

	if vmUUID == "" {
		return "", fmt.Errorf("unable to determine machine ID from cloud configuration or sysfs")
	}

	// Create context
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Create a new finder
	f := find.NewFinder(client.Client, true)

	// Fetch and set data center
	dc, err := f.Datacenter(ctx, cfg.Global.Datacenter)
	if err != nil {
		return "", err
	}
	f.SetDatacenter(dc)

	s := object.NewSearchIndex(client.Client)

	svm, err := s.FindByUuid(ctx, dc, strings.ToLower(strings.TrimSpace(vmUUID)), true, nil)
	if err != nil {
		return "", err
	}

	var vm mo.VirtualMachine
	err = s.Properties(ctx, svm.Reference(), []string{"name"}, &vm)
	if err != nil {
		return "", err
	}
	if log.GetLevel() >= log.DebugLevel {
		log.Debugf("getVMName: vm.Name=%s", vm.Name)
	}
	return vm.Name, nil
}

func newVSphere(cfg VSphereConfig) (*VSphere, error) {

	if cfg.Global.WorkingDir != "" {
		cfg.Global.WorkingDir = path.Clean(cfg.Global.WorkingDir) + "/"
	}
	if cfg.Global.RoundTripperCount == 0 {
		cfg.Global.RoundTripperCount = RoundTripperDefaultCount
	}
	if cfg.Global.VCenterPort != "" {
		log.Warningf("port is a deprecated field in vsphere.conf and will be removed in future release.")
	}

	c, err := newClient(context.TODO(), &cfg)
	if err != nil {
		return nil, err
	}

	id, err := getVMName(c, &cfg)
	if err != nil {
		return nil, err
	}

	vs := VSphere{
		client:          c,
		cfg:             &cfg,
		localInstanceID: id,
	}
	runtime.SetFinalizer(&vs, logout)

	if log.GetLevel() >= log.DebugLevel {
		log.Debugf("newVSphere: vs: %#v", vs)
	}

	return &vs, nil
}

func logout(vs *VSphere) {
	vs.client.Logout(context.TODO())
}

func newClient(ctx context.Context, cfg *VSphereConfig) (*govmomi.Client, error) {
	// Parse URL from string
	u, err := url.Parse(fmt.Sprintf("https://%s/sdk", cfg.Global.VCenterIP))
	if err != nil {
		return nil, err
	}
	// set username and password for the URL
	u.User = url.UserPassword(cfg.Global.User, cfg.Global.Password)

	// Connect and log in to ESX or vCenter
	c, err := govmomi.NewClient(ctx, u, cfg.Global.InsecureFlag)
	if err != nil {
		return nil, err
	}

	// Add retry functionality
	c.RoundTripper = vim25.Retry(c.RoundTripper, vim25.TemporaryNetworkError(int(cfg.Global.RoundTripperCount)))

	return c, nil
}

// Returns a client which communicates with vCenter.
// This client can used to perform further vCenter operations.
func vSphereLogin(ctx context.Context, vs *VSphere) error {
	var err error
	clientLock.Lock()
	defer clientLock.Unlock()
	if vs.client == nil {
		vs.client, err = newClient(ctx, vs.cfg)
		if err != nil {
			return err
		}
		return nil
	}

	m := session.NewManager(vs.client.Client)
	// retrieve client's current session
	u, err := m.UserSession(ctx)
	if err != nil {
		log.Errorf("Error while obtaining user session. err: %q", err)
		return err
	}
	if u != nil {
		return nil
	}

	log.Warningf("Creating new client session since the existing session is not valid or not authenticated")
	vs.client.Logout(ctx)
	vs.client, err = newClient(ctx, vs.cfg)
	if err != nil {
		return err
	}

	return nil
}

// Returns vSphere object `virtual machine` by its name.
func getVirtualMachineByName(ctx context.Context, cfg *VSphereConfig, c *govmomi.Client, nodeName string) (*object.VirtualMachine, error) {
	name := nodeNameToVMName(nodeName)

	if log.GetLevel() >= log.DebugLevel {
		log.Debugf("getVirtualMachineByName: name=%s", name)
	}

	// Create a new finder
	f := find.NewFinder(c.Client, true)

	// Fetch and set data center
	dc, err := f.Datacenter(ctx, cfg.Global.Datacenter)
	if err != nil {
		return nil, err
	}
	f.SetDatacenter(dc)

	vmRegex := cfg.Global.WorkingDir + name

	// Retrieve vm by name
	//TODO: also look for vm inside subfolders
	vm, err := f.VirtualMachine(ctx, vmRegex)
	if err != nil {
		return nil, err
	}

	return vm, nil
}

func getVirtualMachineManagedObjectReference(ctx context.Context, c *govmomi.Client, vm *object.VirtualMachine, field string, dst interface{}) error {
	collector := property.DefaultCollector(c.Client)

	// Retrieve required field from VM object
	err := collector.RetrieveOne(ctx, vm.Reference(), []string{field}, dst)
	if err != nil {
		return err
	}
	return nil
}

// Returns names of running VMs inside VM folder.
func getInstances(ctx context.Context, cfg *VSphereConfig, c *govmomi.Client, filter string) ([]string, error) {
	if log.GetLevel() >= log.DebugLevel {
		log.Debugf("getInstances: filter=%s", filter)
	}

	f := find.NewFinder(c.Client, true)
	dc, err := f.Datacenter(ctx, cfg.Global.Datacenter)
	if err != nil {
		return nil, err
	}

	f.SetDatacenter(dc)

	vmRegex := cfg.Global.WorkingDir + filter

	//TODO: get all vms inside subfolders
	vms, err := f.VirtualMachineList(ctx, vmRegex)
	if err != nil {
		return nil, err
	}

	var vmRef []types.ManagedObjectReference
	for _, vm := range vms {
		vmRef = append(vmRef, vm.Reference())
	}

	pc := property.DefaultCollector(c.Client)

	var vmt []mo.VirtualMachine
	err = pc.Retrieve(ctx, vmRef, []string{"name", "summary"}, &vmt)
	if err != nil {
		return nil, err
	}

	var vmList []string
	for _, vm := range vmt {
		if vm.Summary.Runtime.PowerState == ActivePowerState {
			vmList = append(vmList, vm.Name)
		} else if vm.Summary.Config.Template == false {
			log.Warningf("VM %s, is not in %s state", vm.Name, ActivePowerState)
		}
	}
	return vmList, nil
}

// List returns names of VMs (inside vm folder) by applying filter and which are currently running.
func (vs *VSphere) list(filter string) ([]string, error) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	vmList, err := getInstances(ctx, vs.cfg, vs.client, filter)
	if err != nil {
		return nil, err
	}

	if log.GetLevel() >= log.DebugLevel {
		log.Debugf("Found %d instances matching %s:[ %s ]",
			len(vmList), filter, strings.Join(vmList, ", "))
	}

	var nodeNames []string
	for _, n := range vmList {
		nodeNames = append(nodeNames, n)
	}
	return nodeNames, nil
}

func (vs *VSphere) getAddresses(nodeName string) ([]string, error) {
	addrs := []string{}

	// Create context
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	vm, err := getVirtualMachineByName(ctx, vs.cfg, vs.client, nodeName)
	if err != nil {
		return nil, err
	}

	var mvm mo.VirtualMachine
	err = getVirtualMachineManagedObjectReference(ctx, vs.client, vm, "guest.net", &mvm)
	if err != nil {
		return nil, err
	}

	// retrieve VM's ip(s)
	for _, v := range mvm.Guest.Net {
		if v.Network == vs.cfg.Network.PublicNetwork {
			for _, ip := range v.IpAddress {
				addrs = append(addrs, ip)
			}
		} else if log.GetLevel() >= log.DebugLevel {
			log.Debugf("getAddresses: nodeName=%s, net %v are not in configured network", nodeName, v.IpAddress)
		}
	}
	return addrs, nil
}

// nodeNameToVMName maps a NodeName to the vmware infrastructure name
func nodeNameToVMName(nodeName string) string {
	return string(nodeName)
}
<commit_msg>local client on new vsphere as well<commit_after>package vsphere

import (
	"fmt"
	"io"
	"io/ioutil"
	"net/url"
	"path"
	"runtime"
	"strings"
	"sync"
	"time"

	"gopkg.in/gcfg.v1"

	log "github.com/Sirupsen/logrus"
	"github.com/matt-deboer/etcdcd/pkg/platform"
	"github.com/vmware/govmomi"
	"github.com/vmware/govmomi/find"
	"github.com/vmware/govmomi/object"
	"github.com/vmware/govmomi/property"
	"github.com/vmware/govmomi/session"
	"github.com/vmware/govmomi/vim25"
	"github.com/vmware/govmomi/vim25/mo"
	"github.com/vmware/govmomi/vim25/types"
	"golang.org/x/net/context"

	etcd "github.com/coreos/etcd/client"
)

const (
	ActivePowerState         = "poweredOn"
	RoundTripperDefaultCount = 3
)

var clientLock sync.Mutex

// VSphere is an implementation of cloud provider Interface for VSphere.
type VSphere struct {
	client *govmomi.Client
	cfg    *VSphereConfig
	// InstanceID of the server where this VSphere object is instantiated.
	localInstanceID string
}

type VSphereConfig struct {
	Global struct {
		// vCenter username.
		User string `gcfg:"user"`
		// vCenter password in clear text.
		Password string `gcfg:"password"`
		// vCenter IP.
		VCenterIP string `gcfg:"server"`
		// vCenter port.
		VCenterPort string `gcfg:"port"`
		// True if vCenter uses self-signed cert.
		InsecureFlag bool `gcfg:"insecure-flag"`
		// Datacenter in which VMs are located.
		Datacenter string `gcfg:"datacenter"`
		// Datastore in which vmdks are stored.
		Datastore string `gcfg:"datastore"`
		// WorkingDir is path where VMs can be found.
		WorkingDir string `gcfg:"working-dir"`
		// Soap round tripper count (retries = RoundTripper - 1)
		RoundTripperCount uint `gcfg:"soap-roundtrip-count"`
		// VMUUID is the VM Instance UUID of virtual machine which can be retrieved from instanceUuid
		// property in VmConfigInfo, or also set as vc.uuid in VMX file.
		// If not set, will be fetched from the machine via sysfs (requires root)
		VMUUID string `gcfg:"vm-uuid"`
	}
	Network struct {
		// PublicNetwork is name of the network the VMs are joined to.
		PublicNetwork string `gcfg:"public-network"`
	}
}

func init() {
	platform.Register("vsphere", func(config io.Reader) (platform.Platform, error) {
		cfg, err := readConfig(config)
		if err != nil && !strings.Contains(err.Error(), "warnings") {
			log.Fatal("Failed reading config: ", err)
		}
		return newVSphere(cfg)
	})
}

// Parses vSphere cloud config file and stores it into VSphereConfig.
func readConfig(config io.Reader) (VSphereConfig, error) {
	if config == nil {
		err := fmt.Errorf("no config file given")
		return VSphereConfig{}, err
	}

	var cfg VSphereConfig
	err := gcfg.ReadInto(&cfg, config)
	return cfg, err
}

// ExpectedMembers returns a list of members that should form the cluster
func (vs *VSphere) ExpectedMembers(
	memberFilter string, clientScheme string, clientPort int, serverScheme string, serverPort int) ([]etcd.Member, error) {

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	err := vSphereLogin(ctx, vs)
	if err != nil {
		log.Errorf("VCenter login failed; %v", err)
		return nil, err
	}

	members := []etcd.Member{}
	names, err := vs.list(memberFilter)
	if err != nil {
		return nil, err
	}
	for _, name := range names {
		member := etcd.Member{Name: name, ClientURLs: []string{}, PeerURLs: []string{}}
		for tries := 0; tries <= 10 && len(member.PeerURLs) == 0; tries++ {
			addrs, err := vs.getAddresses(name)
			if err != nil {
				return nil, err
			}
			for _, a := range addrs {
				addr := a
				if strings.Contains(a, ":") {
					addr = "[" + a + "]"
				}
				member.ClientURLs = append(member.ClientURLs, fmt.Sprintf("%s://%s:%d", clientScheme, addr, clientPort))
				member.PeerURLs = append(member.PeerURLs, fmt.Sprintf("%s://%s:%d", serverScheme, addr, serverPort))
			}

			if len(member.PeerURLs) > 0 {
				if log.GetLevel() >= log.DebugLevel {
					log.Debugf("ExpectedMembers: member: %#v", member)
				}
			} else {
				sleepTime := (2 * time.Second)
				if log.GetLevel() >= log.DebugLevel {
					log.Debugf("%s has no addresses yet; sleeping for %s", name, sleepTime)
				}
				time.Sleep(sleepTime)
			}
		}
		members = append(members, member)
	}
	return members, nil
}

// LocalInstanceName returns a list of members that should form the cluster
func (vs *VSphere) LocalInstanceName() string {
	return vs.localInstanceID
}

// Returns the name of the VM on which this code is running.
// Prerequisite: this code assumes VMWare vmtools or open-vm-tools to be installed in the VM.
// Will attempt to determine the machine's name via it's UUID in this precedence order, failing if neither have a UUID:
// * cloud config value VMUUID
// * sysfs entry
func getVMName(client *govmomi.Client, cfg *VSphereConfig) (string, error) {

	var vmUUID string

	if cfg.Global.VMUUID != "" {
		vmUUID = cfg.Global.VMUUID
	} else {
		// This needs root privileges on the host, and will fail otherwise.
		vmUUIDbytes, err := ioutil.ReadFile("/sys/devices/virtual/dmi/id/product_uuid")
		if err != nil {
			return "", err
		}

		vmUUID = string(vmUUIDbytes)
		cfg.Global.VMUUID = vmUUID
	}

	if vmUUID == "" {
		return "", fmt.Errorf("unable to determine machine ID from cloud configuration or sysfs")
	}

	// Create context
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Create a new finder
	f := find.NewFinder(client.Client, true)

	// Fetch and set data center
	dc, err := f.Datacenter(ctx, cfg.Global.Datacenter)
	if err != nil {
		return "", err
	}
	f.SetDatacenter(dc)

	s := object.NewSearchIndex(client.Client)

	svm, err := s.FindByUuid(ctx, dc, strings.ToLower(strings.TrimSpace(vmUUID)), true, nil)
	if err != nil {
		return "", err
	}

	var vm mo.VirtualMachine
	err = s.Properties(ctx, svm.Reference(), []string{"name"}, &vm)
	if err != nil {
		return "", err
	}
	if log.GetLevel() >= log.DebugLevel {
		log.Debugf("getVMName: vm.Name=%s", vm.Name)
	}
	return vm.Name, nil
}

func newVSphere(cfg VSphereConfig) (*VSphere, error) {
	clientLock.Lock()
	defer clientLock.Unlock()

	if cfg.Global.WorkingDir != "" {
		cfg.Global.WorkingDir = path.Clean(cfg.Global.WorkingDir) + "/"
	}
	if cfg.Global.RoundTripperCount == 0 {
		cfg.Global.RoundTripperCount = RoundTripperDefaultCount
	}
	if cfg.Global.VCenterPort != "" {
		log.Warningf("port is a deprecated field in vsphere.conf and will be removed in future release.")
	}

	c, err := newClient(context.TODO(), &cfg)
	if err != nil {
		return nil, err
	}

	id, err := getVMName(c, &cfg)
	if err != nil {
		return nil, err
	}

	vs := VSphere{
		client:          c,
		cfg:             &cfg,
		localInstanceID: id,
	}
	runtime.SetFinalizer(&vs, logout)

	if log.GetLevel() >= log.DebugLevel {
		log.Debugf("newVSphere: vs: %#v", vs)
	}

	return &vs, nil
}

func logout(vs *VSphere) {
	vs.client.Logout(context.TODO())
}

func newClient(ctx context.Context, cfg *VSphereConfig) (*govmomi.Client, error) {
	// Parse URL from string
	u, err := url.Parse(fmt.Sprintf("https://%s/sdk", cfg.Global.VCenterIP))
	if err != nil {
		return nil, err
	}
	// set username and password for the URL
	u.User = url.UserPassword(cfg.Global.User, cfg.Global.Password)

	// Connect and log in to ESX or vCenter
	c, err := govmomi.NewClient(ctx, u, cfg.Global.InsecureFlag)
	if err != nil {
		return nil, err
	}

	// Add retry functionality
	c.RoundTripper = vim25.Retry(c.RoundTripper, vim25.TemporaryNetworkError(int(cfg.Global.RoundTripperCount)))

	return c, nil
}

// Returns a client which communicates with vCenter.
// This client can used to perform further vCenter operations.
func vSphereLogin(ctx context.Context, vs *VSphere) error {
	var err error
	clientLock.Lock()
	defer clientLock.Unlock()
	if vs.client == nil {
		vs.client, err = newClient(ctx, vs.cfg)
		if err != nil {
			return err
		}
		return nil
	}

	m := session.NewManager(vs.client.Client)
	// retrieve client's current session
	u, err := m.UserSession(ctx)
	if err != nil {
		log.Errorf("Error while obtaining user session. err: %q", err)
		return err
	}
	if u != nil {
		return nil
	}

	log.Warningf("Creating new client session since the existing session is not valid or not authenticated")
	vs.client.Logout(ctx)
	vs.client, err = newClient(ctx, vs.cfg)
	if err != nil {
		return err
	}

	return nil
}

// Returns vSphere object `virtual machine` by its name.
func getVirtualMachineByName(ctx context.Context, cfg *VSphereConfig, c *govmomi.Client, nodeName string) (*object.VirtualMachine, error) {
	name := nodeNameToVMName(nodeName)

	if log.GetLevel() >= log.DebugLevel {
		log.Debugf("getVirtualMachineByName: name=%s", name)
	}

	// Create a new finder
	f := find.NewFinder(c.Client, true)

	// Fetch and set data center
	dc, err := f.Datacenter(ctx, cfg.Global.Datacenter)
	if err != nil {
		return nil, err
	}
	f.SetDatacenter(dc)

	vmRegex := cfg.Global.WorkingDir + name

	// Retrieve vm by name
	//TODO: also look for vm inside subfolders
	vm, err := f.VirtualMachine(ctx, vmRegex)
	if err != nil {
		return nil, err
	}

	return vm, nil
}

func getVirtualMachineManagedObjectReference(ctx context.Context, c *govmomi.Client, vm *object.VirtualMachine, field string, dst interface{}) error {
	collector := property.DefaultCollector(c.Client)

	// Retrieve required field from VM object
	err := collector.RetrieveOne(ctx, vm.Reference(), []string{field}, dst)
	if err != nil {
		return err
	}
	return nil
}

// Returns names of running VMs inside VM folder.
func getInstances(ctx context.Context, cfg *VSphereConfig, c *govmomi.Client, filter string) ([]string, error) {
	if log.GetLevel() >= log.DebugLevel {
		log.Debugf("getInstances: filter=%s", filter)
	}

	f := find.NewFinder(c.Client, true)
	dc, err := f.Datacenter(ctx, cfg.Global.Datacenter)
	if err != nil {
		return nil, err
	}

	f.SetDatacenter(dc)

	vmRegex := cfg.Global.WorkingDir + filter

	//TODO: get all vms inside subfolders
	vms, err := f.VirtualMachineList(ctx, vmRegex)
	if err != nil {
		return nil, err
	}

	var vmRef []types.ManagedObjectReference
	for _, vm := range vms {
		vmRef = append(vmRef, vm.Reference())
	}

	pc := property.DefaultCollector(c.Client)

	var vmt []mo.VirtualMachine
	err = pc.Retrieve(ctx, vmRef, []string{"name", "summary"}, &vmt)
	if err != nil {
		return nil, err
	}

	var vmList []string
	for _, vm := range vmt {
		if vm.Summary.Runtime.PowerState == ActivePowerState {
			vmList = append(vmList, vm.Name)
		} else if vm.Summary.Config.Template == false {
			log.Warningf("VM %s, is not in %s state", vm.Name, ActivePowerState)
		}
	}
	return vmList, nil
}

// List returns names of VMs (inside vm folder) by applying filter and which are currently running.
func (vs *VSphere) list(filter string) ([]string, error) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	vmList, err := getInstances(ctx, vs.cfg, vs.client, filter)
	if err != nil {
		return nil, err
	}

	if log.GetLevel() >= log.DebugLevel {
		log.Debugf("Found %d instances matching %s:[ %s ]",
			len(vmList), filter, strings.Join(vmList, ", "))
	}

	var nodeNames []string
	for _, n := range vmList {
		nodeNames = append(nodeNames, n)
	}
	return nodeNames, nil
}

func (vs *VSphere) getAddresses(nodeName string) ([]string, error) {
	addrs := []string{}

	// Create context
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	vm, err := getVirtualMachineByName(ctx, vs.cfg, vs.client, nodeName)
	if err != nil {
		return nil, err
	}

	var mvm mo.VirtualMachine
	err = getVirtualMachineManagedObjectReference(ctx, vs.client, vm, "guest.net", &mvm)
	if err != nil {
		return nil, err
	}

	// retrieve VM's ip(s)
	for _, v := range mvm.Guest.Net {
		if v.Network == vs.cfg.Network.PublicNetwork {
			for _, ip := range v.IpAddress {
				addrs = append(addrs, ip)
			}
		} else if log.GetLevel() >= log.DebugLevel {
			log.Debugf("getAddresses: nodeName=%s, net %v are not in configured network", nodeName, v.IpAddress)
		}
	}
	return addrs, nil
}

// nodeNameToVMName maps a NodeName to the vmware infrastructure name
func nodeNameToVMName(nodeName string) string {
	return string(nodeName)
}
<|endoftext|>


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

// Handlers for pod updates.  Pods map to opflex endpoints

package hostagent

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"strings"

	"github.com/Sirupsen/logrus"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/kubernetes"
	v1 "k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/tools/cache"

	"k8s.io/kubernetes/pkg/controller"

	"github.com/noironetworks/aci-containers/pkg/metadata"
)

type opflexEndpoint struct {
	Uuid string `json:"uuid"`

	EgPolicySpace string                 `json:"eg-policy-space,omitempty"`
	EndpointGroup string                 `json:"endpoint-group-name,omitempty"`
	SecurityGroup []metadata.OpflexGroup `json:"security-group,omitempty"`

	IpAddress  []string `json:"ip,omitempty"`
	MacAddress string   `json:"mac,omitempty"`

	AccessIface       string `json:"access-interface,omitempty"`
	AccessUplinkIface string `json:"access-uplink-interface,omitempty"`
	IfaceName         string `json:"interface-name,omitempty"`

	Attributes map[string]string `json:"attributes,omitempty"`
}

func (agent *HostAgent) initPodInformerFromClient(
	kubeClient *kubernetes.Clientset) {

	agent.initPodInformerBase(
		&cache.ListWatch{
			ListFunc: func(options metav1.ListOptions) (runtime.Object, error) {
				options.FieldSelector =
					fields.Set{"spec.nodeName": agent.config.NodeName}.String()
				return kubeClient.Core().Pods(metav1.NamespaceAll).List(options)
			},
			WatchFunc: func(options metav1.ListOptions) (watch.Interface, error) {
				options.FieldSelector =
					fields.Set{"spec.nodeName": agent.config.NodeName}.String()
				return kubeClient.Core().Pods(metav1.NamespaceAll).Watch(options)
			},
		})
}

func (agent *HostAgent) initPodInformerBase(listWatch *cache.ListWatch) {
	agent.podInformer = cache.NewSharedIndexInformer(
		listWatch,
		&v1.Pod{},
		controller.NoResyncPeriodFunc(),
		cache.Indexers{cache.NamespaceIndex: cache.MetaNamespaceIndexFunc},
	)
	agent.podInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc: func(obj interface{}) {
			agent.podUpdated(obj)
		},
		UpdateFunc: func(_ interface{}, obj interface{}) {
			agent.podUpdated(obj)
		},
		DeleteFunc: func(obj interface{}) {
			agent.podDeleted(obj)
		},
	})
}

func getEp(epfile string) (string, error) {
	raw, err := ioutil.ReadFile(epfile)
	if err != nil {
		return "", err
	}
	return string(raw), err
}

func writeEp(epfile string, ep *opflexEndpoint) (bool, error) {
	newdata, err := json.MarshalIndent(ep, "", "  ")
	if err != nil {
		return true, err
	}
	existingdata, err := ioutil.ReadFile(epfile)
	if err == nil && reflect.DeepEqual(existingdata, newdata) {
		return false, nil
	}

	err = ioutil.WriteFile(epfile, newdata, 0644)
	return true, err
}

func podLogger(log *logrus.Logger, pod *v1.Pod) *logrus.Entry {
	return log.WithFields(logrus.Fields{
		"namespace": pod.ObjectMeta.Namespace,
		"name":      pod.ObjectMeta.Name,
		"node":      pod.Spec.NodeName,
	})
}

func opflexEpLogger(log *logrus.Logger, ep *opflexEndpoint) *logrus.Entry {
	return log.WithFields(logrus.Fields{
		"Uuid":      ep.Uuid,
		"name":      ep.Attributes["vm-name"],
		"namespace": ep.Attributes["namespace"],
	})
}

func (agent *HostAgent) syncEps() {
	if !agent.syncEnabled {
		return
	}

	agent.log.Debug("Syncing endpoints")
	files, err := ioutil.ReadDir(agent.config.OpFlexEndpointDir)
	if err != nil {
		agent.log.WithFields(
			logrus.Fields{"endpointDir": agent.config.OpFlexEndpointDir},
		).Error("Could not read directory ", err)
		return
	}
	seen := make(map[string]bool)
	for _, f := range files {
		if !strings.HasSuffix(f.Name(), ".ep") {
			continue
		}
		epfile := filepath.Join(agent.config.OpFlexEndpointDir, f.Name())
		epidstr := f.Name()
		epidstr = epidstr[:len(epidstr)-3]
		epid := strings.Split(epidstr, "_")

		if len(epid) < 3 {
			agent.log.Warn("Removing invalid endpoint:", f.Name())
			os.Remove(epfile)
			continue
		}
		poduuid := epid[0]
		contid := epid[1]
		contiface := epid[2]

		logger := agent.log.WithFields(
			logrus.Fields{
				"PodUuid":   poduuid,
				"ContId":    contid,
				"ContIFace": contiface,
			},
		)

		existing, ok := agent.opflexEps[poduuid]
		if ok {
			ok = false
			for _, ep := range existing {
				if ep.Uuid != epidstr {
					continue
				}

				wrote, err := writeEp(epfile, ep)
				if err != nil {
					opflexEpLogger(agent.log, ep).
						Error("Error writing EP file: ", err)
				} else if wrote {
					opflexEpLogger(agent.log, ep).
						Info("Updated endpoint")
				}
				seen[epidstr] = true
				ok = true
			}
		}
		if !ok {
			logger.Info("Removing endpoint")
			os.Remove(epfile)
		}
	}

	for _, eps := range agent.opflexEps {
		for _, ep := range eps {
			if seen[ep.Uuid] {
				continue
			}

			opflexEpLogger(agent.log, ep).Info("Adding endpoint")
			epfile := filepath.Join(agent.config.OpFlexEndpointDir,
				ep.Uuid+".ep")
			_, err = writeEp(epfile, ep)
			if err != nil {
				opflexEpLogger(agent.log, ep).
					Error("Error writing EP file: ", err)
			}
		}
	}
	agent.log.Debug("Finished endpoint sync")
}

func podFilter(pod *v1.Pod) bool {
	if pod.Spec.HostNetwork {
		return false
	}
	return true
}

func (agent *HostAgent) podUpdated(obj interface{}) {
	agent.indexMutex.Lock()
	defer agent.indexMutex.Unlock()
	agent.podChangedLocked(obj)
}

func (agent *HostAgent) podChanged(podkey *string) {
	podobj, exists, err := agent.podInformer.GetStore().GetByKey(*podkey)
	if err != nil {
		agent.log.Error("Could not lookup pod: ", err)
	}
	if !exists || podobj == nil {
		agent.log.Info("Object doesn't exist yet ", *podkey)
		return
	}

	agent.indexMutex.Lock()
	defer agent.indexMutex.Unlock()
	agent.podChangedLocked(podobj)
}

func (agent *HostAgent) podChangedLocked(podobj interface{}) {
	pod := podobj.(*v1.Pod)
	logger := podLogger(agent.log, pod)

	if !podFilter(pod) {
		delete(agent.opflexEps, string(pod.ObjectMeta.UID))
		agent.syncEps()
		return
	}

	id := fmt.Sprintf("%s/%s", pod.ObjectMeta.Namespace, pod.ObjectMeta.Name)
	epmetadata, ok := agent.epMetadata[id]
	if !ok {
		logger.Debug("No metadata")
		delete(agent.opflexEps, string(pod.ObjectMeta.UID))
		agent.syncEps()
		return
	}

	var neweps []*opflexEndpoint

	for _, epmeta := range epmetadata {
		for _, iface := range epmeta.Ifaces {
			patchIntName, patchAccessName :=
				metadata.GetIfaceNames(iface.HostVethName)

			ips := make([]string, 0)
			for _, ip := range iface.IPs {
				if ip.Address.IP == nil {
					continue
				}
				ips = append(ips, ip.Address.IP.String())
			}

			epidstr := string(pod.ObjectMeta.UID) + "_" +
				epmeta.Id.ContId + "_" + iface.HostVethName
			ep := &opflexEndpoint{
				Uuid:              epidstr,
				MacAddress:        iface.Mac,
				IpAddress:         ips,
				AccessIface:       iface.HostVethName,
				AccessUplinkIface: patchAccessName,
				IfaceName:         patchIntName,
			}

			ep.Attributes = pod.ObjectMeta.Labels
			ep.Attributes["vm-name"] = pod.ObjectMeta.Name
			ep.Attributes["namespace"] = pod.ObjectMeta.Namespace
			ep.Attributes["interface-name"] = iface.HostVethName

			if egval, ok := pod.ObjectMeta.Annotations[metadata.CompEgAnnotation]; ok {
				g := &metadata.OpflexGroup{}
				err := json.Unmarshal([]byte(egval), g)
				if err != nil {
					logger.WithFields(logrus.Fields{
						"EgAnnotation": egval,
					}).Error("Could not decode annotation: ", err)
				} else {
					ep.EgPolicySpace = g.PolicySpace
					ep.EndpointGroup = g.Name
				}
			}
			if sgval, ok := pod.ObjectMeta.Annotations[metadata.CompSgAnnotation]; ok {
				g := make([]metadata.OpflexGroup, 0)
				err := json.Unmarshal([]byte(sgval), &g)
				if err != nil {
					logger.WithFields(logrus.Fields{
						"SgAnnotation": sgval,
					}).Error("Could not decode annotation: ", err)
				} else {
					ep.SecurityGroup = g
				}
			}

			neweps = append(neweps, ep)
		}
	}

	existing, ok := agent.opflexEps[string(pod.ObjectMeta.UID)]
	if (ok && !reflect.DeepEqual(existing, neweps)) || !ok {
		logger.WithFields(logrus.Fields{
			"id": id,
			"ep": neweps,
		}).Debug("Updated endpoints for pod")

		agent.opflexEps[string(pod.ObjectMeta.UID)] = neweps

		agent.syncEps()
	}
}

func (agent *HostAgent) podDeleted(obj interface{}) {
	agent.indexMutex.Lock()
	defer agent.indexMutex.Unlock()

	agent.podDeletedLocked(obj)
}

func (agent *HostAgent) podDeletedLocked(obj interface{}) {
	pod := obj.(*v1.Pod)
	u := string(pod.ObjectMeta.UID)
	if _, ok := agent.opflexEps[u]; ok {
		delete(agent.opflexEps, u)
		agent.syncEps()
	}
}
<commit_msg>Allow generating EP attributed even when pod has no labels.<commit_after>// Copyright 2016 Cisco Systems, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Handlers for pod updates.  Pods map to opflex endpoints

package hostagent

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"strings"

	"github.com/Sirupsen/logrus"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/kubernetes"
	v1 "k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/tools/cache"

	"k8s.io/kubernetes/pkg/controller"

	"github.com/noironetworks/aci-containers/pkg/metadata"
)

type opflexEndpoint struct {
	Uuid string `json:"uuid"`

	EgPolicySpace string                 `json:"eg-policy-space,omitempty"`
	EndpointGroup string                 `json:"endpoint-group-name,omitempty"`
	SecurityGroup []metadata.OpflexGroup `json:"security-group,omitempty"`

	IpAddress  []string `json:"ip,omitempty"`
	MacAddress string   `json:"mac,omitempty"`

	AccessIface       string `json:"access-interface,omitempty"`
	AccessUplinkIface string `json:"access-uplink-interface,omitempty"`
	IfaceName         string `json:"interface-name,omitempty"`

	Attributes map[string]string `json:"attributes,omitempty"`
}

func (agent *HostAgent) initPodInformerFromClient(
	kubeClient *kubernetes.Clientset) {

	agent.initPodInformerBase(
		&cache.ListWatch{
			ListFunc: func(options metav1.ListOptions) (runtime.Object, error) {
				options.FieldSelector =
					fields.Set{"spec.nodeName": agent.config.NodeName}.String()
				return kubeClient.Core().Pods(metav1.NamespaceAll).List(options)
			},
			WatchFunc: func(options metav1.ListOptions) (watch.Interface, error) {
				options.FieldSelector =
					fields.Set{"spec.nodeName": agent.config.NodeName}.String()
				return kubeClient.Core().Pods(metav1.NamespaceAll).Watch(options)
			},
		})
}

func (agent *HostAgent) initPodInformerBase(listWatch *cache.ListWatch) {
	agent.podInformer = cache.NewSharedIndexInformer(
		listWatch,
		&v1.Pod{},
		controller.NoResyncPeriodFunc(),
		cache.Indexers{cache.NamespaceIndex: cache.MetaNamespaceIndexFunc},
	)
	agent.podInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc: func(obj interface{}) {
			agent.podUpdated(obj)
		},
		UpdateFunc: func(_ interface{}, obj interface{}) {
			agent.podUpdated(obj)
		},
		DeleteFunc: func(obj interface{}) {
			agent.podDeleted(obj)
		},
	})
}

func getEp(epfile string) (string, error) {
	raw, err := ioutil.ReadFile(epfile)
	if err != nil {
		return "", err
	}
	return string(raw), err
}

func writeEp(epfile string, ep *opflexEndpoint) (bool, error) {
	newdata, err := json.MarshalIndent(ep, "", "  ")
	if err != nil {
		return true, err
	}
	existingdata, err := ioutil.ReadFile(epfile)
	if err == nil && reflect.DeepEqual(existingdata, newdata) {
		return false, nil
	}

	err = ioutil.WriteFile(epfile, newdata, 0644)
	return true, err
}

func podLogger(log *logrus.Logger, pod *v1.Pod) *logrus.Entry {
	return log.WithFields(logrus.Fields{
		"namespace": pod.ObjectMeta.Namespace,
		"name":      pod.ObjectMeta.Name,
		"node":      pod.Spec.NodeName,
	})
}

func opflexEpLogger(log *logrus.Logger, ep *opflexEndpoint) *logrus.Entry {
	return log.WithFields(logrus.Fields{
		"Uuid":      ep.Uuid,
		"name":      ep.Attributes["vm-name"],
		"namespace": ep.Attributes["namespace"],
	})
}

func (agent *HostAgent) syncEps() {
	if !agent.syncEnabled {
		return
	}

	agent.log.Debug("Syncing endpoints")
	files, err := ioutil.ReadDir(agent.config.OpFlexEndpointDir)
	if err != nil {
		agent.log.WithFields(
			logrus.Fields{"endpointDir": agent.config.OpFlexEndpointDir},
		).Error("Could not read directory ", err)
		return
	}
	seen := make(map[string]bool)
	for _, f := range files {
		if !strings.HasSuffix(f.Name(), ".ep") {
			continue
		}
		epfile := filepath.Join(agent.config.OpFlexEndpointDir, f.Name())
		epidstr := f.Name()
		epidstr = epidstr[:len(epidstr)-3]
		epid := strings.Split(epidstr, "_")

		if len(epid) < 3 {
			agent.log.Warn("Removing invalid endpoint:", f.Name())
			os.Remove(epfile)
			continue
		}
		poduuid := epid[0]
		contid := epid[1]
		contiface := epid[2]

		logger := agent.log.WithFields(
			logrus.Fields{
				"PodUuid":   poduuid,
				"ContId":    contid,
				"ContIFace": contiface,
			},
		)

		existing, ok := agent.opflexEps[poduuid]
		if ok {
			ok = false
			for _, ep := range existing {
				if ep.Uuid != epidstr {
					continue
				}

				wrote, err := writeEp(epfile, ep)
				if err != nil {
					opflexEpLogger(agent.log, ep).
						Error("Error writing EP file: ", err)
				} else if wrote {
					opflexEpLogger(agent.log, ep).
						Info("Updated endpoint")
				}
				seen[epidstr] = true
				ok = true
			}
		}
		if !ok {
			logger.Info("Removing endpoint")
			os.Remove(epfile)
		}
	}

	for _, eps := range agent.opflexEps {
		for _, ep := range eps {
			if seen[ep.Uuid] {
				continue
			}

			opflexEpLogger(agent.log, ep).Info("Adding endpoint")
			epfile := filepath.Join(agent.config.OpFlexEndpointDir,
				ep.Uuid+".ep")
			_, err = writeEp(epfile, ep)
			if err != nil {
				opflexEpLogger(agent.log, ep).
					Error("Error writing EP file: ", err)
			}
		}
	}
	agent.log.Debug("Finished endpoint sync")
}

func podFilter(pod *v1.Pod) bool {
	if pod.Spec.HostNetwork {
		return false
	}
	return true
}

func (agent *HostAgent) podUpdated(obj interface{}) {
	agent.indexMutex.Lock()
	defer agent.indexMutex.Unlock()
	agent.podChangedLocked(obj)
}

func (agent *HostAgent) podChanged(podkey *string) {
	podobj, exists, err := agent.podInformer.GetStore().GetByKey(*podkey)
	if err != nil {
		agent.log.Error("Could not lookup pod: ", err)
	}
	if !exists || podobj == nil {
		agent.log.Info("Object doesn't exist yet ", *podkey)
		return
	}

	agent.indexMutex.Lock()
	defer agent.indexMutex.Unlock()
	agent.podChangedLocked(podobj)
}

func (agent *HostAgent) podChangedLocked(podobj interface{}) {
	pod := podobj.(*v1.Pod)
	logger := podLogger(agent.log, pod)

	if !podFilter(pod) {
		delete(agent.opflexEps, string(pod.ObjectMeta.UID))
		agent.syncEps()
		return
	}

	id := fmt.Sprintf("%s/%s", pod.ObjectMeta.Namespace, pod.ObjectMeta.Name)
	epmetadata, ok := agent.epMetadata[id]
	if !ok {
		logger.Debug("No metadata")
		delete(agent.opflexEps, string(pod.ObjectMeta.UID))
		agent.syncEps()
		return
	}

	var neweps []*opflexEndpoint

	for _, epmeta := range epmetadata {
		for _, iface := range epmeta.Ifaces {
			patchIntName, patchAccessName :=
				metadata.GetIfaceNames(iface.HostVethName)

			ips := make([]string, 0)
			for _, ip := range iface.IPs {
				if ip.Address.IP == nil {
					continue
				}
				ips = append(ips, ip.Address.IP.String())
			}

			epidstr := string(pod.ObjectMeta.UID) + "_" +
				epmeta.Id.ContId + "_" + iface.HostVethName
			ep := &opflexEndpoint{
				Uuid:              epidstr,
				MacAddress:        iface.Mac,
				IpAddress:         ips,
				AccessIface:       iface.HostVethName,
				AccessUplinkIface: patchAccessName,
				IfaceName:         patchIntName,
			}

			ep.Attributes = pod.ObjectMeta.Labels
			if ep.Attributes == nil {
				ep.Attributes = make(map[string]string)
			}
			ep.Attributes["vm-name"] = pod.ObjectMeta.Name
			ep.Attributes["namespace"] = pod.ObjectMeta.Namespace
			ep.Attributes["interface-name"] = iface.HostVethName

			if egval, ok := pod.ObjectMeta.Annotations[metadata.CompEgAnnotation]; ok {
				g := &metadata.OpflexGroup{}
				err := json.Unmarshal([]byte(egval), g)
				if err != nil {
					logger.WithFields(logrus.Fields{
						"EgAnnotation": egval,
					}).Error("Could not decode annotation: ", err)
				} else {
					ep.EgPolicySpace = g.PolicySpace
					ep.EndpointGroup = g.Name
				}
			}
			if sgval, ok := pod.ObjectMeta.Annotations[metadata.CompSgAnnotation]; ok {
				g := make([]metadata.OpflexGroup, 0)
				err := json.Unmarshal([]byte(sgval), &g)
				if err != nil {
					logger.WithFields(logrus.Fields{
						"SgAnnotation": sgval,
					}).Error("Could not decode annotation: ", err)
				} else {
					ep.SecurityGroup = g
				}
			}

			neweps = append(neweps, ep)
		}
	}

	existing, ok := agent.opflexEps[string(pod.ObjectMeta.UID)]
	if (ok && !reflect.DeepEqual(existing, neweps)) || !ok {
		logger.WithFields(logrus.Fields{
			"id": id,
			"ep": neweps,
		}).Debug("Updated endpoints for pod")

		agent.opflexEps[string(pod.ObjectMeta.UID)] = neweps

		agent.syncEps()
	}
}

func (agent *HostAgent) podDeleted(obj interface{}) {
	agent.indexMutex.Lock()
	defer agent.indexMutex.Unlock()

	agent.podDeletedLocked(obj)
}

func (agent *HostAgent) podDeletedLocked(obj interface{}) {
	pod := obj.(*v1.Pod)
	u := string(pod.ObjectMeta.UID)
	if _, ok := agent.opflexEps[u]; ok {
		delete(agent.opflexEps, u)
		agent.syncEps()
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The cert-manager 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 scheduler

import (
	"runtime"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"k8s.io/utils/clock"
)

func Test_afterFunc(t *testing.T) {
	// Note that re-implimenting AfterFunc is not a good idea, since testing it
	// is tricky as seen in time_test.go in the standard library. We will just
	// focus on two important cases: "f" should be run after the duration

	t.Run("stop works", func(t *testing.T) {
		// This test makes sure afterFunc does not leak goroutines.
		//
		// This test may be run concurrently to other tests, so we want to avoid
		// being affected by the goroutines from other tests. To do that, we start a
		// huge number of afterFunc and check that the number of goroutines before
		// and after more or less match.
		expected := runtime.NumGoroutine()
		var cancels []func()
		for i := 1; i <= 10000; i++ {
			cancels = append(cancels, afterFunc(clock.RealClock{}, 1*time.Hour, func() {
				t.Errorf("should never be called")
			}))
		}

		for _, cancel := range cancels {
			cancel()
		}

		// We don't know when the goroutines will actually finish.
		time.Sleep(100 * time.Millisecond)

		t.Logf("%d goroutines before, %d goroutines after", expected, runtime.NumGoroutine())

		assert.InDelta(t, expected, runtime.NumGoroutine(), 100)
	})

	t.Run("f is called after 100 milliseconds", func(t *testing.T) {
		var end time.Time
		wait := make(chan struct{})

		start := time.Now()

		afterFunc(clock.RealClock{}, 100*time.Millisecond, func() {
			end = time.Now()
			close(wait)
		})

		<-wait
		assert.InDelta(t, 100*time.Millisecond, end.Sub(start), float64(1*time.Millisecond))
	})
}

func TestAdd(t *testing.T) {
	after := newMockAfter()

	var wg sync.WaitGroup
	type testT struct {
		obj      string
		duration time.Duration
	}
	tests := []testT{
		{"test500", time.Millisecond * 500},
		{"test1000", time.Second * 1},
		{"test3000", time.Second * 3},
	}
	for _, test := range tests {
		wg.Add(1)
		t.Run(test.obj, func(test testT) func(*testing.T) {
			waitSubtest := make(chan struct{})
			return func(t *testing.T) {
				startTime := after.currentTime
				queue := NewScheduledWorkQueue(clock.RealClock{}, func(obj interface{}) {
					defer wg.Done()
					durationEarly := test.duration - after.currentTime.Sub(startTime)

					if durationEarly > 0 {
						t.Errorf("got queue item %.2f seconds too early", float64(durationEarly)/float64(time.Second))
					}
					if obj != test.obj {
						t.Errorf("expected obj '%+v' but got obj '%+v'", test.obj, obj)
					}
					waitSubtest <- struct{}{}
				})
				queue.(*scheduledWorkQueue).afterFunc = after.AfterFunc
				queue.Add(test.obj, test.duration)
				after.warp(test.duration + time.Millisecond)
				<-waitSubtest
			}
		}(test))
	}

	wg.Wait()
}

func TestForget(t *testing.T) {
	after := newMockAfter()

	var wg sync.WaitGroup
	type testT struct {
		obj      string
		duration time.Duration
	}
	tests := []testT{
		{"test500", time.Millisecond * 500},
		{"test1000", time.Second * 1},
		{"test3000", time.Second * 3},
	}
	for _, test := range tests {
		wg.Add(1)
		t.Run(test.obj, func(test testT) func(*testing.T) {
			return func(t *testing.T) {
				defer wg.Done()
				queue := NewScheduledWorkQueue(clock.RealClock{}, func(_ interface{}) {
					t.Errorf("scheduled function should never be called")
				})
				queue.(*scheduledWorkQueue).afterFunc = after.AfterFunc
				queue.Add(test.obj, test.duration)
				queue.Forget(test.obj)
				after.warp(test.duration * 2)
			}
		}(test))
	}

	wg.Wait()
}

// TestConcurrentAdd checks that if we add the same item concurrently, it
// doesn't end up hitting a data-race / leaking a timer.
func TestConcurrentAdd(t *testing.T) {
	after := newMockAfter()

	var wg sync.WaitGroup
	queue := NewScheduledWorkQueue(clock.RealClock{}, func(obj interface{}) {
		t.Fatalf("should not be called, but was called with %v", obj)
	})
	queue.(*scheduledWorkQueue).afterFunc = after.AfterFunc

	for i := 0; i < 1000; i++ {
		wg.Add(1)
		go func() {
			queue.Add(1, 1*time.Second)
			wg.Done()
		}()
	}
	wg.Wait()

	queue.Forget(1)
	after.warp(5 * time.Second)
}

type timerQueueItem struct {
	f       func()
	t       time.Time
	run     bool
	stopped bool
}

func (tq *timerQueueItem) Stop() bool {
	stopped := tq.stopped
	tq.stopped = true
	return stopped
}

type mockAfter struct {
	lock        *sync.Mutex
	currentTime time.Time
	queue       []*timerQueueItem
}

func newMockAfter() *mockAfter {
	return &mockAfter{
		queue: make([]*timerQueueItem, 0),
		lock:  &sync.Mutex{},
	}
}

func (m *mockAfter) AfterFunc(c clock.Clock, d time.Duration, f func()) func() {
	m.lock.Lock()
	defer m.lock.Unlock()

	item := &timerQueueItem{
		f: f,
		t: m.currentTime.Add(d),
	}
	m.queue = append(m.queue, item)
	return func() {
		item.Stop()
	}
}

func (m *mockAfter) warp(d time.Duration) {
	m.lock.Lock()
	defer m.lock.Unlock()
	m.currentTime = m.currentTime.Add(d)
	for _, item := range m.queue {
		if item.run || item.stopped {
			continue
		}

		if item.t.Before(m.currentTime) {
			item.run = true
			go item.f()
		}
	}
}
<commit_msg>Increase margin of error in an otherwise unsound test<commit_after>/*
Copyright 2020 The cert-manager 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 scheduler

import (
	"runtime"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"k8s.io/utils/clock"
)

func Test_afterFunc(t *testing.T) {
	// Note that re-implimenting AfterFunc is not a good idea, since testing it
	// is tricky as seen in time_test.go in the standard library. We will just
	// focus on two important cases: "f" should be run after the duration

	t.Run("stop works", func(t *testing.T) {
		// This test makes sure afterFunc does not leak goroutines.
		//
		// This test may be run concurrently to other tests, so we want to avoid
		// being affected by the goroutines from other tests. To do that, we start a
		// huge number of afterFunc and check that the number of goroutines before
		// and after more or less match.
		expected := runtime.NumGoroutine()
		var cancels []func()
		for i := 1; i <= 10000; i++ {
			cancels = append(cancels, afterFunc(clock.RealClock{}, 1*time.Hour, func() {
				t.Errorf("should never be called")
			}))
		}

		for _, cancel := range cancels {
			cancel()
		}

		// We don't know when the goroutines will actually finish.
		time.Sleep(100 * time.Millisecond)

		t.Logf("%d goroutines before, %d goroutines after", expected, runtime.NumGoroutine())

		assert.InDelta(t, expected, runtime.NumGoroutine(), 100)
	})

	t.Run("f is called after roughly 100 milliseconds", func(t *testing.T) {
		var end time.Time
		wait := make(chan struct{})

		start := time.Now()

		expectedDuration := 100 * time.Millisecond

		afterFunc(clock.RealClock{}, expectedDuration, func() {
			end = time.Now()
			close(wait)
		})

		<-wait

		elapsed := end.Sub(start)

		assert.InDelta(t, expectedDuration, elapsed, float64(500*time.Millisecond))
		assert.GreaterOrEqual(t, elapsed, expectedDuration)
	})
}

func TestAdd(t *testing.T) {
	after := newMockAfter()

	var wg sync.WaitGroup
	type testT struct {
		obj      string
		duration time.Duration
	}
	tests := []testT{
		{"test500", time.Millisecond * 500},
		{"test1000", time.Second * 1},
		{"test3000", time.Second * 3},
	}
	for _, test := range tests {
		wg.Add(1)
		t.Run(test.obj, func(test testT) func(*testing.T) {
			waitSubtest := make(chan struct{})
			return func(t *testing.T) {
				startTime := after.currentTime
				queue := NewScheduledWorkQueue(clock.RealClock{}, func(obj interface{}) {
					defer wg.Done()
					durationEarly := test.duration - after.currentTime.Sub(startTime)

					if durationEarly > 0 {
						t.Errorf("got queue item %.2f seconds too early", float64(durationEarly)/float64(time.Second))
					}
					if obj != test.obj {
						t.Errorf("expected obj '%+v' but got obj '%+v'", test.obj, obj)
					}
					waitSubtest <- struct{}{}
				})
				queue.(*scheduledWorkQueue).afterFunc = after.AfterFunc
				queue.Add(test.obj, test.duration)
				after.warp(test.duration + time.Millisecond)
				<-waitSubtest
			}
		}(test))
	}

	wg.Wait()
}

func TestForget(t *testing.T) {
	after := newMockAfter()

	var wg sync.WaitGroup
	type testT struct {
		obj      string
		duration time.Duration
	}
	tests := []testT{
		{"test500", time.Millisecond * 500},
		{"test1000", time.Second * 1},
		{"test3000", time.Second * 3},
	}
	for _, test := range tests {
		wg.Add(1)
		t.Run(test.obj, func(test testT) func(*testing.T) {
			return func(t *testing.T) {
				defer wg.Done()
				queue := NewScheduledWorkQueue(clock.RealClock{}, func(_ interface{}) {
					t.Errorf("scheduled function should never be called")
				})
				queue.(*scheduledWorkQueue).afterFunc = after.AfterFunc
				queue.Add(test.obj, test.duration)
				queue.Forget(test.obj)
				after.warp(test.duration * 2)
			}
		}(test))
	}

	wg.Wait()
}

// TestConcurrentAdd checks that if we add the same item concurrently, it
// doesn't end up hitting a data-race / leaking a timer.
func TestConcurrentAdd(t *testing.T) {
	after := newMockAfter()

	var wg sync.WaitGroup
	queue := NewScheduledWorkQueue(clock.RealClock{}, func(obj interface{}) {
		t.Fatalf("should not be called, but was called with %v", obj)
	})
	queue.(*scheduledWorkQueue).afterFunc = after.AfterFunc

	for i := 0; i < 1000; i++ {
		wg.Add(1)
		go func() {
			queue.Add(1, 1*time.Second)
			wg.Done()
		}()
	}
	wg.Wait()

	queue.Forget(1)
	after.warp(5 * time.Second)
}

type timerQueueItem struct {
	f       func()
	t       time.Time
	run     bool
	stopped bool
}

func (tq *timerQueueItem) Stop() bool {
	stopped := tq.stopped
	tq.stopped = true
	return stopped
}

type mockAfter struct {
	lock        *sync.Mutex
	currentTime time.Time
	queue       []*timerQueueItem
}

func newMockAfter() *mockAfter {
	return &mockAfter{
		queue: make([]*timerQueueItem, 0),
		lock:  &sync.Mutex{},
	}
}

func (m *mockAfter) AfterFunc(c clock.Clock, d time.Duration, f func()) func() {
	m.lock.Lock()
	defer m.lock.Unlock()

	item := &timerQueueItem{
		f: f,
		t: m.currentTime.Add(d),
	}
	m.queue = append(m.queue, item)
	return func() {
		item.Stop()
	}
}

func (m *mockAfter) warp(d time.Duration) {
	m.lock.Lock()
	defer m.lock.Unlock()
	m.currentTime = m.currentTime.Add(d)
	for _, item := range m.queue {
		if item.run || item.stopped {
			continue
		}

		if item.t.Before(m.currentTime) {
			item.run = true
			go item.f()
		}
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Camlistore 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 nodeattr contains constants for permanode attribute names.
package nodeattr

const (
	// DateCreated is http://schema.org/dateCreated: The date on
	// which an item was created, in RFC 3339 format (with
	// Camlistore's addition that zone -00:01 means localtime:
	// unknown timezone).
	DateCreated = "dateCreated"

	// Title is http://schema.org/title
	Title = "title"

	// Description is http://schema.org/description
	// Value is plain text, no HTML, newlines are newlines.
	Description = "description"

	// Type is the Camlistore permanode type ("camliNodeType").
	// Importer-specific ones are of the form "domain.com:objecttype".
	// Well-defined ones are documented in doc/schema/claims/attributes.txt.
	Type = "camliNodeType"

	// CamliContent is "camliContent", the blobref of the permanode's content.
	// For files or images, the camliContent is fileref (the blobref of
	// the "file" schema blob).
	CamliContent = "camliContent"

	// Content is "content", used e.g. for the content of a tweet.
	// TODO: define this more
	Content = "content"

	// StartDate is http://schema.org/startDate: The start
	// date and time of the event or item (in RFC 3339 date
	// format).
	StartDate = "startDate"

	// URL is the item's original or origin URL.
	URL = "url"

	Latitude  = "latitude"
	Longitude = "longitude"
)
<commit_msg>nodeattr: add more, rearrange, clean up some comments.<commit_after>/*
Copyright 2014 The Camlistore 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 nodeattr contains constants for permanode attribute names.
//
// For all date values in RFC 3339 format, Camlistore additionally
// treats the special timezone offset -00:01 (one minute west of UTC)
// as meaning that the local time was known, but the location or
// timezone was not. Usually this is from EXIF files.
package nodeattr

const (
	// Type is the Camlistore permanode type ("camliNodeType").
	// Importer-specific ones are of the form "domain.com:objecttype".
	// Well-defined ones are documented in doc/schema/claims/attributes.txt.
	Type = "camliNodeType"

	// CamliContent is "camliContent", the blobref of the permanode's content.
	// For files or images, the camliContent is fileref (the blobref of
	// the "file" schema blob).
	CamliContent = "camliContent"

	// DateCreated is http://schema.org/dateCreated in RFC 3339
	// format.
	DateCreated = "dateCreated"

	// StartDate is http://schema.org/startDate, the start date
	// and time of the event or item, in RFC 3339 format.
	StartDate = "startDate"

	// DateModified is http://schema.org/dateModified, in RFC 3339
	// format.
	DateModified = "dateModified"

	// DatePublished is http://schema.org/datePublished in RFC
	// 3339 format.
	DatePublished = "datePublished"

	// Title is http://schema.org/title
	Title = "title"

	// Description is http://schema.org/description
	// Value is plain text, no HTML, newlines are newlines.
	Description = "description"

	// Content is "content", used e.g. for the content of a tweet.
	// TODO: define this more
	Content = "content"

	// URL is the item's original or origin URL.
	URL = "url"

	// LocationText is free-flowing text definition of a location or place, such
	// as a city name, or a full postal address.
	LocationText = "locationText"

	Latitude  = "latitude"
	Longitude = "longitude"
)
<|endoftext|>


<commit_before>// Copyright 2015-present Oursky Ltd.
//
// 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 zmq

package zmq

import (
	"bytes"
	"errors"
	"fmt"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/zeromq/goczmq"
)

const (
	// HeartbeatInterval is the interval that broker and worker send
	// heartbeats to each other.
	HeartbeatInterval = time.Second
	// HeartbeatLiveness defines the liveness of each heartbeat. Generally
	// it should be >= 3, otherwise workers will keep being discarded and
	// reconnecting.
	HeartbeatLiveness = 3
)

var heartbeatIntervalMS = int(HeartbeatInterval.Seconds() * 1000)

const (
	// Ready is sent by worker to signal broker that it is ready to receive
	// jobs.
	Ready = "\001"
	// Heartbeat is sent by both broker and worker to signal a heartbeat.
	Heartbeat = "\002"
	// Shutdown is sent by worker while being killed (probably by CTRL C).
	// It is an addition to original PPP to shorten the time needed for
	// broker to detect a normal shutdown of worker.
	Shutdown = "\003"
)

// Broker implements the Paranoid Pirate queue described in the zguide:
// http://zguide.zeromq.org/py:all#Robust-Reliable-Queuing-Paranoid-Pirate-Pattern
// Related RFC: https://rfc.zeromq.org/spec:6/PPP
// with the addition of:
//
// 1. Shutdown signal, which signifies a normal termination of worker to provide
//    a fast path of worker removal
type Broker struct {
	name string
	// NOTE: goroutines are caller of plugin, so frontend is Go side,
	// backend is plugin side
	frontend      *goczmq.Sock
	backend       *goczmq.Sock
	bothPoller    *goczmq.Poller
	backendPoller *goczmq.Poller
	workers       workerQueue
	freshWorkers  chan []byte
	logger        *logrus.Entry
}

// NewBroker returns a new *Broker.
func NewBroker(name, frontendAddr, backendAddr string) (*Broker, error) {
	namedLogger := log.WithFields(logrus.Fields{"plugin": name})
	frontend, err := goczmq.NewRouter(frontendAddr)
	if err != nil {
		panic(err)
	}

	backend, err := goczmq.NewRouter(backendAddr)
	if err != nil {
		panic(err)
	}

	backendPoller, err := goczmq.NewPoller(backend)
	if err != nil {
		panic(err)
	}

	bothPoller, err := goczmq.NewPoller(frontend, backend)
	if err != nil {
		panic(err)
	}

	return &Broker{
		name:          name,
		frontend:      frontend,
		backend:       backend,
		bothPoller:    bothPoller,
		backendPoller: backendPoller,
		workers:       newWorkerQueue(),
		freshWorkers:  make(chan []byte, 1),
		logger:        namedLogger,
	}, nil
}

// Run kicks start the Broker and listens for requests. It blocks function
// execution.
func (lb *Broker) Run() {
	heartbeatAt := time.Now().Add(HeartbeatInterval)
	for {
		var sock *goczmq.Sock
		if lb.workers.Len() == 0 {
			sock = lb.backendPoller.Wait(heartbeatIntervalMS)
		} else {
			sock = lb.bothPoller.Wait(heartbeatIntervalMS)
		}

		switch sock {
		case lb.backend:
			frames, err := lb.backend.RecvMessage()
			if err != nil {
				panic(err)
			}

			address := frames[0]
			msg := frames[1:]
			tErr := lb.workers.Tick(newWorker(address))
			if tErr != nil {
				status := string(msg[0])
				if status != Ready {
					lb.logger.Warnln(tErr)
				}
			}
			if len(msg) == 1 {
				status := string(msg[0])
				lb.handleWorkerStatus(&lb.workers, address, status)
			} else {
				lb.frontend.SendMessage(msg)
				lb.logger.Debugf("zmq/broker: plugin => server: %#x, %s\n", msg[0], msg)
			}
		case lb.frontend:
			frames, err := lb.frontend.RecvMessage()
			if err != nil {
				panic(err)
			}

			frames = append([][]byte{lb.workers.Next()}, frames...)
			lb.backend.SendMessage(frames)
			lb.logger.Debugf("zmq/broker: server => plugin: %#x, %s\n", frames[0], frames)
		case nil:
			// do nothing
		default:
			panic("zmq/broker: received unknown socket")
		}

		lb.logger.Debugf("zmq/broker: idle worker count %d\n", lb.workers.Len())
		if heartbeatAt.Before(time.Now()) {
			for _, worker := range lb.workers.pworkers {
				msg := [][]byte{worker.address, []byte(Heartbeat)}
				lb.logger.Debugf("zmq/broker: server => plugin Heartbeat: %s\n", worker.address)
				lb.backend.SendMessage(msg)
			}
			heartbeatAt = time.Now().Add(HeartbeatInterval)
		}

		lb.workers.Purge()
	}
}

func (lb *Broker) handleWorkerStatus(workers *workerQueue, address []byte, status string) {
	switch status {
	case Ready:
		log.Infof("zmq/broker: ready worker = %s", address)
		workers.Add(newWorker(address))
		lb.freshWorkers <- address
	case Heartbeat:
		// no-ops
	case Shutdown:
		workers.Remove(address)
		log.Infof("zmq/broker: shutdown of worker = %s", address)
	default:
		log.Errorf("zmq/broker: invalid status from worker = %s: %s", address, status)
	}
}

type pworker struct {
	address []byte
	expiry  time.Time
}

func newWorker(address []byte) pworker {
	return pworker{
		address,
		time.Now().Add(HeartbeatLiveness * HeartbeatInterval),
	}
}

// workerQueue is a last tick fist out queue.
// A worker need to register itself using Add before it can tick.
// Ticking of an non-registered worker will be no-ops.
type workerQueue struct {
	pworkers  []pworker
	addresses map[string]bool
}

func newWorkerQueue() workerQueue {
	return workerQueue{
		[]pworker{},
		map[string]bool{},
	}
}

func (q workerQueue) Len() int {
	return len(q.pworkers)
}

func (q *workerQueue) Next() []byte {
	workers := q.pworkers
	worker := workers[len(workers)-1]
	q.pworkers = workers[:len(workers)-1]
	return worker.address
}

func (q *workerQueue) Add(worker pworker) {
	q.addresses[string(worker.address)] = true
	err := q.Tick(worker)
	if err == nil {
		return
	}
}

func (q *workerQueue) Tick(worker pworker) error {
	if _, ok := q.addresses[string(worker.address)]; !ok {
		return errors.New(fmt.Sprintf("zmq/broker: Ticking non-registered worker = %s", worker.address))
	}
	workers := q.pworkers

	for i, w := range workers {
		if bytes.Equal(w.address, worker.address) {
			q.pworkers = append(append(workers[:i], workers[i+1:]...), worker)
			return nil
		}
	}
	q.pworkers = append(q.pworkers, worker)
	log.Debugf("zmq/broker: worker return to poll = %s", worker.address)
	return nil
}

func (q *workerQueue) Purge() {
	workers := q.pworkers

	now := time.Now()
	for i, w := range workers {
		if w.expiry.After(now) {
			break
		}
		q.pworkers = workers[i+1:]
		delete(q.addresses, string(w.address))
		log.Infof("zmq/broker: disconnected worker = %s", w.address)
	}
}

func (q *workerQueue) Remove(address []byte) {
	delete(q.addresses, string(address))
	workers := q.pworkers

	for i, w := range workers {
		if bytes.Equal(w.address, address) {
			q.pworkers = append(workers[:i], workers[i+1:]...)
			break
		}
	}
}
<commit_msg>Fix a minor typo<commit_after>// Copyright 2015-present Oursky Ltd.
//
// 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 zmq

package zmq

import (
	"bytes"
	"errors"
	"fmt"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/zeromq/goczmq"
)

const (
	// HeartbeatInterval is the interval that broker and worker send
	// heartbeats to each other.
	HeartbeatInterval = time.Second
	// HeartbeatLiveness defines the liveness of each heartbeat. Generally
	// it should be >= 3, otherwise workers will keep being discarded and
	// reconnecting.
	HeartbeatLiveness = 3
)

var heartbeatIntervalMS = int(HeartbeatInterval.Seconds() * 1000)

const (
	// Ready is sent by worker to signal broker that it is ready to receive
	// jobs.
	Ready = "\001"
	// Heartbeat is sent by both broker and worker to signal a heartbeat.
	Heartbeat = "\002"
	// Shutdown is sent by worker while being killed (probably by CTRL C).
	// It is an addition to original PPP to shorten the time needed for
	// broker to detect a normal shutdown of worker.
	Shutdown = "\003"
)

// Broker implements the Paranoid Pirate queue described in the zguide:
// http://zguide.zeromq.org/py:all#Robust-Reliable-Queuing-Paranoid-Pirate-Pattern
// Related RFC: https://rfc.zeromq.org/spec:6/PPP
// with the addition of:
//
// 1. Shutdown signal, which signifies a normal termination of worker to provide
//    a fast path of worker removal
type Broker struct {
	name string
	// NOTE: goroutines are caller of plugin, so frontend is Go side,
	// backend is plugin side
	frontend      *goczmq.Sock
	backend       *goczmq.Sock
	bothPoller    *goczmq.Poller
	backendPoller *goczmq.Poller
	workers       workerQueue
	freshWorkers  chan []byte
	logger        *logrus.Entry
}

// NewBroker returns a new *Broker.
func NewBroker(name, frontendAddr, backendAddr string) (*Broker, error) {
	namedLogger := log.WithFields(logrus.Fields{"plugin": name})
	frontend, err := goczmq.NewRouter(frontendAddr)
	if err != nil {
		panic(err)
	}

	backend, err := goczmq.NewRouter(backendAddr)
	if err != nil {
		panic(err)
	}

	backendPoller, err := goczmq.NewPoller(backend)
	if err != nil {
		panic(err)
	}

	bothPoller, err := goczmq.NewPoller(frontend, backend)
	if err != nil {
		panic(err)
	}

	return &Broker{
		name:          name,
		frontend:      frontend,
		backend:       backend,
		bothPoller:    bothPoller,
		backendPoller: backendPoller,
		workers:       newWorkerQueue(),
		freshWorkers:  make(chan []byte, 1),
		logger:        namedLogger,
	}, nil
}

// Run kicks start the Broker and listens for requests. It blocks function
// execution.
func (lb *Broker) Run() {
	heartbeatAt := time.Now().Add(HeartbeatInterval)
	for {
		var sock *goczmq.Sock
		if lb.workers.Len() == 0 {
			sock = lb.backendPoller.Wait(heartbeatIntervalMS)
		} else {
			sock = lb.bothPoller.Wait(heartbeatIntervalMS)
		}

		switch sock {
		case lb.backend:
			frames, err := lb.backend.RecvMessage()
			if err != nil {
				panic(err)
			}

			address := frames[0]
			msg := frames[1:]
			tErr := lb.workers.Tick(newWorker(address))
			if tErr != nil {
				status := string(msg[0])
				if status != Ready {
					lb.logger.Warnln(tErr)
				}
			}
			if len(msg) == 1 {
				status := string(msg[0])
				lb.handleWorkerStatus(&lb.workers, address, status)
			} else {
				lb.frontend.SendMessage(msg)
				lb.logger.Debugf("zmq/broker: plugin => server: %#x, %s\n", msg[0], msg)
			}
		case lb.frontend:
			frames, err := lb.frontend.RecvMessage()
			if err != nil {
				panic(err)
			}

			frames = append([][]byte{lb.workers.Next()}, frames...)
			lb.backend.SendMessage(frames)
			lb.logger.Debugf("zmq/broker: server => plugin: %#x, %s\n", frames[0], frames)
		case nil:
			// do nothing
		default:
			panic("zmq/broker: received unknown socket")
		}

		lb.logger.Debugf("zmq/broker: idle worker count %d\n", lb.workers.Len())
		if heartbeatAt.Before(time.Now()) {
			for _, worker := range lb.workers.pworkers {
				msg := [][]byte{worker.address, []byte(Heartbeat)}
				lb.logger.Debugf("zmq/broker: server => plugin Heartbeat: %s\n", worker.address)
				lb.backend.SendMessage(msg)
			}
			heartbeatAt = time.Now().Add(HeartbeatInterval)
		}

		lb.workers.Purge()
	}
}

func (lb *Broker) handleWorkerStatus(workers *workerQueue, address []byte, status string) {
	switch status {
	case Ready:
		log.Infof("zmq/broker: ready worker = %s", address)
		workers.Add(newWorker(address))
		lb.freshWorkers <- address
	case Heartbeat:
		// no-op
	case Shutdown:
		workers.Remove(address)
		log.Infof("zmq/broker: shutdown of worker = %s", address)
	default:
		log.Errorf("zmq/broker: invalid status from worker = %s: %s", address, status)
	}
}

type pworker struct {
	address []byte
	expiry  time.Time
}

func newWorker(address []byte) pworker {
	return pworker{
		address,
		time.Now().Add(HeartbeatLiveness * HeartbeatInterval),
	}
}

// workerQueue is a last tick fist out queue.
// A worker need to register itself using Add before it can tick.
// Ticking of an non-registered worker will be no-ops.
type workerQueue struct {
	pworkers  []pworker
	addresses map[string]bool
}

func newWorkerQueue() workerQueue {
	return workerQueue{
		[]pworker{},
		map[string]bool{},
	}
}

func (q workerQueue) Len() int {
	return len(q.pworkers)
}

func (q *workerQueue) Next() []byte {
	workers := q.pworkers
	worker := workers[len(workers)-1]
	q.pworkers = workers[:len(workers)-1]
	return worker.address
}

func (q *workerQueue) Add(worker pworker) {
	q.addresses[string(worker.address)] = true
	err := q.Tick(worker)
	if err == nil {
		return
	}
}

func (q *workerQueue) Tick(worker pworker) error {
	if _, ok := q.addresses[string(worker.address)]; !ok {
		return errors.New(fmt.Sprintf("zmq/broker: Ticking non-registered worker = %s", worker.address))
	}
	workers := q.pworkers

	for i, w := range workers {
		if bytes.Equal(w.address, worker.address) {
			q.pworkers = append(append(workers[:i], workers[i+1:]...), worker)
			return nil
		}
	}
	q.pworkers = append(q.pworkers, worker)
	log.Debugf("zmq/broker: worker return to poll = %s", worker.address)
	return nil
}

func (q *workerQueue) Purge() {
	workers := q.pworkers

	now := time.Now()
	for i, w := range workers {
		if w.expiry.After(now) {
			break
		}
		q.pworkers = workers[i+1:]
		delete(q.addresses, string(w.address))
		log.Infof("zmq/broker: disconnected worker = %s", w.address)
	}
}

func (q *workerQueue) Remove(address []byte) {
	delete(q.addresses, string(address))
	workers := q.pworkers

	for i, w := range workers {
		if bytes.Equal(w.address, address) {
			q.pworkers = append(workers[:i], workers[i+1:]...)
			break
		}
	}
}
<|endoftext|>


<commit_before>package alerting

import (
	"context"
	"time"

	"github.com/benbjohnson/clock"
	"github.com/grafana/grafana/pkg/log"
	"golang.org/x/sync/errgroup"
)

type Engine struct {
	execQueue     chan *Job
	clock         clock.Clock
	ticker        *Ticker
	scheduler     Scheduler
	evalHandler   EvalHandler
	ruleReader    RuleReader
	log           log.Logger
	resultHandler ResultHandler
}

func NewEngine() *Engine {
	e := &Engine{
		ticker:        NewTicker(time.Now(), time.Second*0, clock.New()),
		execQueue:     make(chan *Job, 1000),
		scheduler:     NewScheduler(),
		evalHandler:   NewEvalHandler(),
		ruleReader:    NewRuleReader(),
		log:           log.New("alerting.engine"),
		resultHandler: NewResultHandler(),
	}

	return e
}

func (e *Engine) Run(ctx context.Context) error {
	e.log.Info("Initializing Alerting")

	alertGroup, ctx := errgroup.WithContext(ctx)

	alertGroup.Go(func() error { return e.alertingTicker(ctx) })
	alertGroup.Go(func() error { return e.runJobDispatcher(ctx) })

	err := alertGroup.Wait()

	e.log.Info("Stopped Alerting", "reason", err)
	return err
}

func (e *Engine) alertingTicker(grafanaCtx context.Context) error {
	defer func() {
		if err := recover(); err != nil {
			e.log.Error("Scheduler Panic: stopping alertingTicker", "error", err, "stack", log.Stack(1))
		}
	}()

	tickIndex := 0

	for {
		select {
		case <-grafanaCtx.Done():
			return grafanaCtx.Err()
		case tick := <-e.ticker.C:
			// TEMP SOLUTION update rules ever tenth tick
			if tickIndex%10 == 0 {
				e.scheduler.Update(e.ruleReader.Fetch())
			}

			e.scheduler.Tick(tick, e.execQueue)
			tickIndex++
		}
	}
}

func (e *Engine) runJobDispatcher(grafanaCtx context.Context) error {
	dispatcherGroup, alertCtx := errgroup.WithContext(grafanaCtx)

	for {
		select {
		case <-grafanaCtx.Done():
			return dispatcherGroup.Wait()
		case job := <-e.execQueue:
			dispatcherGroup.Go(func() error { return e.processJob(alertCtx, job) })
		}
	}
}

var (
	unfinishedWorkTimeout time.Duration = time.Second * 5
	alertTimeout          time.Duration = time.Second * 30
)

func (e *Engine) processJob(grafanaCtx context.Context, job *Job) error {
	defer func() {
		if err := recover(); err != nil {
			e.log.Error("Alert Panic", "error", err, "stack", log.Stack(1))
		}
	}()

	alertCtx, cancelFn := context.WithTimeout(context.TODO(), alertTimeout)

	job.Running = true
	evalContext := NewEvalContext(alertCtx, job.Rule)

	done := make(chan struct{})

	go func() {
		e.evalHandler.Eval(evalContext)
		e.resultHandler.Handle(evalContext)
		close(done)
	}()

	var err error = nil
	select {
	case <-grafanaCtx.Done():
		select {
		case <-time.After(unfinishedWorkTimeout):
			cancelFn()
			err = grafanaCtx.Err()
		case <-done:
		}
	case <-done:
	}

	e.log.Debug("Job Execution completed", "timeMs", evalContext.GetDurationMs(), "alertId", evalContext.Rule.Id, "name", evalContext.Rule.Name, "firing", evalContext.Firing)
	job.Running = false
	cancelFn()
	return err
}
<commit_msg>feat(alerting): recover from panic<commit_after>package alerting

import (
	"context"
	"time"

	"github.com/benbjohnson/clock"
	"github.com/grafana/grafana/pkg/log"
	"golang.org/x/sync/errgroup"
)

type Engine struct {
	execQueue     chan *Job
	clock         clock.Clock
	ticker        *Ticker
	scheduler     Scheduler
	evalHandler   EvalHandler
	ruleReader    RuleReader
	log           log.Logger
	resultHandler ResultHandler
}

func NewEngine() *Engine {
	e := &Engine{
		ticker:        NewTicker(time.Now(), time.Second*0, clock.New()),
		execQueue:     make(chan *Job, 1000),
		scheduler:     NewScheduler(),
		evalHandler:   NewEvalHandler(),
		ruleReader:    NewRuleReader(),
		log:           log.New("alerting.engine"),
		resultHandler: NewResultHandler(),
	}

	return e
}

func (e *Engine) Run(ctx context.Context) error {
	e.log.Info("Initializing Alerting")

	alertGroup, ctx := errgroup.WithContext(ctx)

	alertGroup.Go(func() error { return e.alertingTicker(ctx) })
	alertGroup.Go(func() error { return e.runJobDispatcher(ctx) })

	err := alertGroup.Wait()

	e.log.Info("Stopped Alerting", "reason", err)
	return err
}

func (e *Engine) alertingTicker(grafanaCtx context.Context) error {
	defer func() {
		if err := recover(); err != nil {
			e.log.Error("Scheduler Panic: stopping alertingTicker", "error", err, "stack", log.Stack(1))
		}
	}()

	tickIndex := 0

	for {
		select {
		case <-grafanaCtx.Done():
			return grafanaCtx.Err()
		case tick := <-e.ticker.C:
			// TEMP SOLUTION update rules ever tenth tick
			if tickIndex%10 == 0 {
				e.scheduler.Update(e.ruleReader.Fetch())
			}

			e.scheduler.Tick(tick, e.execQueue)
			tickIndex++
		}
	}
}

func (e *Engine) runJobDispatcher(grafanaCtx context.Context) error {
	dispatcherGroup, alertCtx := errgroup.WithContext(grafanaCtx)

	for {
		select {
		case <-grafanaCtx.Done():
			return dispatcherGroup.Wait()
		case job := <-e.execQueue:
			dispatcherGroup.Go(func() error { return e.processJob(alertCtx, job) })
		}
	}
}

var (
	unfinishedWorkTimeout time.Duration = time.Second * 5
	alertTimeout          time.Duration = time.Second * 30
)

func (e *Engine) processJob(grafanaCtx context.Context, job *Job) error {
	defer func() {
		if err := recover(); err != nil {
			e.log.Error("Alert Panic", "error", err, "stack", log.Stack(1))
		}
	}()

	alertCtx, cancelFn := context.WithTimeout(context.TODO(), alertTimeout)

	job.Running = true
	evalContext := NewEvalContext(alertCtx, job.Rule)

	done := make(chan struct{})

	go func() {
		defer func() {
			if err := recover(); err != nil {
				e.log.Error("Alert Panic", "error", err, "stack", log.Stack(1))
				close(done)
			}
		}()

		e.evalHandler.Eval(evalContext)
		e.resultHandler.Handle(evalContext)
		close(done)
	}()

	var err error = nil
	select {
	case <-grafanaCtx.Done():
		select {
		case <-time.After(unfinishedWorkTimeout):
			cancelFn()
			err = grafanaCtx.Err()
		case <-done:
		}
	case <-done:
	}

	e.log.Debug("Job Execution completed", "timeMs", evalContext.GetDurationMs(), "alertId", evalContext.Rule.Id, "name", evalContext.Rule.Name, "firing", evalContext.Firing)
	job.Running = false
	cancelFn()
	return err
}
<|endoftext|>


<commit_before>package sqlstore

import (
	"github.com/go-xorm/xorm"
	"github.com/grafana/grafana/pkg/bus"
	"github.com/grafana/grafana/pkg/log"
)

type dbTransactionFunc func(sess *xorm.Session) error
type dbTransactionFunc2 func(sess *session) error

type session struct {
	*xorm.Session
	events []interface{}
}

func (sess *session) publishAfterCommit(msg interface{}) {
	sess.events = append(sess.events, msg)
}

func inTransaction(callback dbTransactionFunc) error {
	var err error

	sess := x.NewSession()
	defer sess.Close()

	if err = sess.Begin(); err != nil {
		return err
	}

	err = callback(sess)

	if err != nil {
		sess.Rollback()
		return err
	} else if err = sess.Commit(); err != nil {
		return err
	}

	return nil
}

func inTransaction2(callback dbTransactionFunc2) error {
	var err error

	sess := session{Session: x.NewSession()}

	defer sess.Close()
	if err = sess.Begin(); err != nil {
		return err
	}

	err = callback(&sess)

	if err != nil {
		sess.Rollback()
		return err
	} else if err = sess.Commit(); err != nil {
		return err
	}

	if len(sess.events) > 0 {
		for _, e := range sess.events {
			if err = bus.Publish(e); err != nil {
				log.Error(3, "Failed to publish event after commit", err)
			}
		}
	}

	return nil
}
<commit_msg>sqlite: fixed database table looked handling, now retries up to 5 times, fixes #7992<commit_after>package sqlstore

import (
	"time"

	"github.com/go-xorm/xorm"
	"github.com/grafana/grafana/pkg/bus"
	"github.com/grafana/grafana/pkg/log"
	sqlite3 "github.com/mattn/go-sqlite3"
)

type dbTransactionFunc func(sess *xorm.Session) error
type dbTransactionFunc2 func(sess *session) error

type session struct {
	*xorm.Session
	events []interface{}
}

func (sess *session) publishAfterCommit(msg interface{}) {
	sess.events = append(sess.events, msg)
}

func inTransaction(callback dbTransactionFunc) error {
	return inTransactionWithRetry(callback, 0)
}

func inTransactionWithRetry(callback dbTransactionFunc, retry int) error {
	var err error

	sess := x.NewSession()
	defer sess.Close()

	if err = sess.Begin(); err != nil {
		return err
	}

	err = callback(sess)

	// special handling of database locked errors for sqlite, then we can retry 3 times
	if sqlError, ok := err.(sqlite3.Error); ok && retry < 5 {
		if sqlError.Code == sqlite3.ErrLocked {
			sess.Rollback()
			time.Sleep(time.Millisecond * time.Duration(10))
			sqlog.Info("Database table locked, sleeping then retrying", "retry", retry)
			return inTransactionWithRetry(callback, retry+1)
		}
	}

	if err != nil {
		sess.Rollback()
		return err
	} else if err = sess.Commit(); err != nil {
		return err
	}

	return nil
}

func inTransaction2(callback dbTransactionFunc2) error {
	var err error

	sess := session{Session: x.NewSession()}

	defer sess.Close()
	if err = sess.Begin(); err != nil {
		return err
	}

	err = callback(&sess)

	if err != nil {
		sess.Rollback()
		return err
	} else if err = sess.Commit(); err != nil {
		return err
	}

	if len(sess.events) > 0 {
		for _, e := range sess.events {
			if err = bus.Publish(e); err != nil {
				log.Error(3, "Failed to publish event after commit", err)
			}
		}
	}

	return nil
}
<|endoftext|>