<commit_before>package cluster

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"regexp"

	"github.com/rancher/norman/types/convert"
	v3 "github.com/rancher/rancher/pkg/apis/management.cattle.io/v3"
	v1 "github.com/rancher/rancher/pkg/apis/provisioning.cattle.io/v1"
	"github.com/rancher/rancher/pkg/features"
	capicontrollers "github.com/rancher/rancher/pkg/generated/controllers/cluster.x-k8s.io/v1alpha4"
	mgmtcontrollers "github.com/rancher/rancher/pkg/generated/controllers/management.cattle.io/v3"
	rocontrollers "github.com/rancher/rancher/pkg/generated/controllers/provisioning.cattle.io/v1"
	"github.com/rancher/rancher/pkg/provisioningv2/kubeconfig"
	"github.com/rancher/rancher/pkg/wrangler"
	"github.com/rancher/wrangler/pkg/apply"
	"github.com/rancher/wrangler/pkg/condition"
	corecontrollers "github.com/rancher/wrangler/pkg/generated/controllers/core/v1"
	"github.com/rancher/wrangler/pkg/generic"
	"github.com/rancher/wrangler/pkg/kstatus"
	"github.com/rancher/wrangler/pkg/name"
	"github.com/rancher/wrangler/pkg/relatedresource"
	"github.com/rancher/wrangler/pkg/yaml"
	rbacv1 "k8s.io/api/rbac/v1"
	apierror "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
)

const (
	ByCluster    = "by-cluster"
	creatorIDAnn = "field.cattle.io/creatorId"
)

var (
	mgmtNameRegexp = regexp.MustCompile("c-[a-z0-9]{5}|local")
)

type handler struct {
	mgmtClusterCache  mgmtcontrollers.ClusterCache
	mgmtClusters      mgmtcontrollers.ClusterClient
	clusterTokenCache mgmtcontrollers.ClusterRegistrationTokenCache
	clusterTokens     mgmtcontrollers.ClusterRegistrationTokenClient
	clusters          rocontrollers.ClusterController
	clusterCache      rocontrollers.ClusterCache
	secretCache       corecontrollers.SecretCache
	kubeconfigManager *kubeconfig.Manager
	apply             apply.Apply

	capiClusters capicontrollers.ClusterCache
	capiMachines capicontrollers.MachineCache
}

func Register(
	ctx context.Context,
	clients *wrangler.Context) {
	h := handler{
		mgmtClusterCache:  clients.Mgmt.Cluster().Cache(),
		mgmtClusters:      clients.Mgmt.Cluster(),
		clusterTokenCache: clients.Mgmt.ClusterRegistrationToken().Cache(),
		clusterTokens:     clients.Mgmt.ClusterRegistrationToken(),
		clusters:          clients.Provisioning.Cluster(),
		clusterCache:      clients.Provisioning.Cluster().Cache(),
		secretCache:       clients.Core.Secret().Cache(),
		capiClusters:      clients.CAPI.Cluster().Cache(),
		capiMachines:      clients.CAPI.Machine().Cache(),
		kubeconfigManager: kubeconfig.New(clients),
		apply: clients.Apply.WithCacheTypes(
			clients.Provisioning.Cluster(),
			clients.Mgmt.Cluster()),
	}

	mgmtcontrollers.RegisterClusterGeneratingHandler(ctx,
		clients.Mgmt.Cluster(),
		clients.Apply.WithCacheTypes(clients.Provisioning.Cluster()),
		"",
		"provisioning-cluster-create",
		h.generateProvisioning,
		nil)

	rocontrollers.RegisterClusterGeneratingHandler(ctx,
		clients.Provisioning.Cluster(),
		clients.Apply.WithCacheTypes(clients.Mgmt.Cluster(),
			clients.Mgmt.ClusterRoleTemplateBinding(),
			clients.Mgmt.ClusterRegistrationToken(),
			clients.RBAC.ClusterRoleBinding(),
			clients.Core.Namespace(),
			clients.Core.Secret()),
		"Created",
		"cluster-create",
		h.generateCluster,
		&generic.GeneratingHandlerOptions{
			AllowClusterScoped: true,
		},
	)

	clients.Mgmt.Cluster().OnChange(ctx, "cluster-watch", h.createToken)
	clusterCache := clients.Provisioning.Cluster().Cache()
	relatedresource.Watch(ctx, "cluster-watch", h.clusterWatch,
		clients.Provisioning.Cluster(), clients.Mgmt.Cluster())

	clusterCache.AddIndexer(ByCluster, byClusterIndex)

	clients.Mgmt.Cluster().OnRemove(ctx, "mgmt-cluster-remove", h.OnMgmtClusterRemove)
	clients.Provisioning.Cluster().OnRemove(ctx, "provisioning-cluster-remove", h.OnClusterRemove)
}

func byClusterIndex(obj *v1.Cluster) ([]string, error) {
	if obj.Status.ClusterName == "" {
		return nil, nil
	}
	return []string{obj.Status.ClusterName}, nil
}

func (h *handler) clusterWatch(namespace, name string, obj runtime.Object) ([]relatedresource.Key, error) {
	cluster, ok := obj.(*v3.Cluster)
	if !ok {
		return nil, nil
	}
	operatorClusters, err := h.clusterCache.GetByIndex(ByCluster, cluster.Name)
	if err != nil || len(operatorClusters) == 0 {
		// ignore
		return nil, nil
	}
	return []relatedresource.Key{
		{
			Namespace: operatorClusters[0].Namespace,
			Name:      operatorClusters[0].Name,
		},
	}, nil
}

func (h *handler) isLegacyCluster(cluster interface{}) bool {
	if c, ok := cluster.(*v3.Cluster); ok {
		return mgmtNameRegexp.MatchString(c.Name)
	} else if c, ok := cluster.(*v1.Cluster); ok {
		return mgmtNameRegexp.MatchString(c.Name)
	}
	return false
}

func (h *handler) generateProvisioning(cluster *v3.Cluster, status v3.ClusterStatus) ([]runtime.Object, v3.ClusterStatus, error) {
	if !h.isLegacyCluster(cluster) || cluster.Spec.FleetWorkspaceName == "" {
		return nil, status, nil
	}
	return []runtime.Object{
		&v1.Cluster{
			ObjectMeta: metav1.ObjectMeta{
				Name:        cluster.Name,
				Namespace:   cluster.Spec.FleetWorkspaceName,
				Labels:      yaml.CleanAnnotationsForExport(cluster.Labels),
				Annotations: yaml.CleanAnnotationsForExport(cluster.Annotations),
			},
		},
	}, status, nil
}

func (h *handler) generateCluster(cluster *v1.Cluster, status v1.ClusterStatus) ([]runtime.Object, v1.ClusterStatus, error) {
	switch {
	case cluster.Spec.ClusterAPIConfig != nil:
		return h.createClusterAndDeployAgent(cluster, status)
	default:
		return h.createCluster(cluster, status, v3.ClusterSpec{
			ImportedConfig: &v3.ImportedConfig{},
		})
	}
}

func NormalizeCluster(cluster *v3.Cluster) (runtime.Object, error) {
	// We do this so that we don't clobber status because the rancher object is pretty dirty and doesn't have a status subresource
	data, err := convert.EncodeToMap(cluster)
	if err != nil {
		return nil, err
	}
	data = map[string]interface{}{
		"metadata": data["metadata"],
		"spec":     data["spec"],
	}
	data["kind"] = "Cluster"
	data["apiVersion"] = "management.cattle.io/v3"
	return &unstructured.Unstructured{Object: data}, nil
}

func (h *handler) createToken(_ string, cluster *v3.Cluster) (*v3.Cluster, error) {
	if cluster == nil {
		return cluster, nil
	}
	_, err := h.clusterTokenCache.Get(cluster.Name, "default-token")
	if apierror.IsNotFound(err) {
		_, err = h.clusterTokens.Create(&v3.ClusterRegistrationToken{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "default-token",
				Namespace: cluster.Name,
			},
			Spec: v3.ClusterRegistrationTokenSpec{
				ClusterName: cluster.Name,
			},
		})
		return cluster, err
	}
	return cluster, err
}

func (h *handler) createCluster(cluster *v1.Cluster, status v1.ClusterStatus, spec v3.ClusterSpec) ([]runtime.Object, v1.ClusterStatus, error) {
	if h.isLegacyCluster(cluster) {
		mgmtCluster, err := h.mgmtClusterCache.Get(cluster.Name)
		if err != nil {
			return nil, status, err
		}
		return h.updateStatus(nil, cluster, status, mgmtCluster)
	}
	return h.createNewCluster(cluster, status, spec)
}

func mgmtClusterName(clusterNamespace, clusterName string) string {
	hash := sha256.Sum256([]byte(clusterNamespace + "/" + clusterName))
	return name.SafeConcatName("c", "m", hex.EncodeToString(hash[:])[:8])
}

func (h *handler) createNewCluster(cluster *v1.Cluster, status v1.ClusterStatus, spec v3.ClusterSpec) ([]runtime.Object, v1.ClusterStatus, error) {
	spec.DisplayName = cluster.Name
	spec.Description = cluster.Annotations["field.cattle.io/description"]
	spec.FleetWorkspaceName = cluster.Namespace
	spec.AgentEnvVars = cluster.Spec.AgentEnvVars
	spec.DefaultPodSecurityPolicyTemplateName = cluster.Spec.DefaultPodSecurityPolicyTemplateName
	spec.DefaultClusterRoleForProjectMembers = cluster.Spec.DefaultClusterRoleForProjectMembers
	spec.EnableNetworkPolicy = cluster.Spec.EnableNetworkPolicy

	if cluster.Spec.RKEConfig != nil {
		spec.LocalClusterAuthEndpoint = v3.LocalClusterAuthEndpoint{
			FQDN:    cluster.Spec.RKEConfig.LocalClusterAuthEndpoint.FQDN,
			CACerts: cluster.Spec.RKEConfig.LocalClusterAuthEndpoint.CACerts,
			Enabled: cluster.Spec.RKEConfig.LocalClusterAuthEndpoint.Enabled,
		}
	}

	newCluster := &v3.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Name:        mgmtClusterName(cluster.Namespace, cluster.Name),
			Labels:      cluster.Labels,
			Annotations: map[string]string{},
		},
		Spec: spec,
	}

	for k, v := range cluster.Annotations {
		newCluster.Annotations[k] = v
	}

	userName, err := h.kubeconfigManager.EnsureUser(cluster.Namespace, cluster.Name)
	if err != nil {
		return nil, status, err
	}

	newCluster.Annotations[creatorIDAnn] = userName

	normalizedCluster, err := NormalizeCluster(newCluster)
	if err != nil {
		return nil, status, err
	}

	return h.updateStatus([]runtime.Object{
		normalizedCluster,
	}, cluster, status, newCluster)
}

func (h *handler) updateStatus(objs []runtime.Object, cluster *v1.Cluster, status v1.ClusterStatus, rCluster *v3.Cluster) ([]runtime.Object, v1.ClusterStatus, error) {
	ready := false
	existing, err := h.mgmtClusterCache.Get(rCluster.Name)
	if err != nil && !apierror.IsNotFound(err) {
		return nil, status, err
	} else if err == nil {
		if condition.Cond("Ready").IsTrue(existing) {
			ready = true
		}
	}

	// Never set ready back to false because we will end up deleting the secret
	status.Ready = status.Ready || ready
	status.ObservedGeneration = cluster.Generation
	status.ClusterName = rCluster.Name
	if ready {
		kstatus.SetActive(&status)
	} else {
		kstatus.SetTransitioning(&status, "")
	}

	if status.Ready {
		secret, err := h.kubeconfigManager.GetKubeConfig(cluster, status)
		if err != nil {
			return nil, status, err
		}
		if secret != nil {
			if secret.UID == "" {
				objs = append(objs, secret)
			}
			status.ClientSecretName = secret.Name

			if features.MCM.Enabled() {
				crtb, err := h.kubeconfigManager.GetCRTBForAdmin(cluster, status)
				if err != nil {
					return nil, status, err
				}
				objs = append(objs, crtb)
			} else if cluster.Namespace == "fleet-local" && cluster.Name == "local" {
				user, err := h.kubeconfigManager.EnsureUser(cluster.Namespace, cluster.Name)
				if err != nil {
					return objs, status, err
				}
				objs = append(objs, &rbacv1.ClusterRoleBinding{
					ObjectMeta: metav1.ObjectMeta{
						Name: "fleet-local-cluster-admin",
					},
					Subjects: []rbacv1.Subject{{
						Kind:     "User",
						APIGroup: rbacv1.GroupName,
						Name:     user,
					}},
					RoleRef: rbacv1.RoleRef{
						Kind: "ClusterRole",
						Name: "cluster-admin",
					},
				})
			}
		}
	}

	return objs, status, nil
}
<commit_msg>Only use crtbs when mcm is enabled<commit_after>package cluster

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"regexp"

	"github.com/rancher/norman/types/convert"
	v3 "github.com/rancher/rancher/pkg/apis/management.cattle.io/v3"
	v1 "github.com/rancher/rancher/pkg/apis/provisioning.cattle.io/v1"
	"github.com/rancher/rancher/pkg/features"
	capicontrollers "github.com/rancher/rancher/pkg/generated/controllers/cluster.x-k8s.io/v1alpha4"
	mgmtcontrollers "github.com/rancher/rancher/pkg/generated/controllers/management.cattle.io/v3"
	rocontrollers "github.com/rancher/rancher/pkg/generated/controllers/provisioning.cattle.io/v1"
	"github.com/rancher/rancher/pkg/provisioningv2/kubeconfig"
	"github.com/rancher/rancher/pkg/wrangler"
	"github.com/rancher/wrangler/pkg/apply"
	"github.com/rancher/wrangler/pkg/condition"
	corecontrollers "github.com/rancher/wrangler/pkg/generated/controllers/core/v1"
	"github.com/rancher/wrangler/pkg/generic"
	"github.com/rancher/wrangler/pkg/kstatus"
	"github.com/rancher/wrangler/pkg/name"
	"github.com/rancher/wrangler/pkg/relatedresource"
	"github.com/rancher/wrangler/pkg/yaml"
	rbacv1 "k8s.io/api/rbac/v1"
	apierror "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
)

const (
	ByCluster    = "by-cluster"
	creatorIDAnn = "field.cattle.io/creatorId"
)

var (
	mgmtNameRegexp = regexp.MustCompile("c-[a-z0-9]{5}|local")
)

type handler struct {
	mgmtClusterCache  mgmtcontrollers.ClusterCache
	mgmtClusters      mgmtcontrollers.ClusterClient
	clusterTokenCache mgmtcontrollers.ClusterRegistrationTokenCache
	clusterTokens     mgmtcontrollers.ClusterRegistrationTokenClient
	clusters          rocontrollers.ClusterController
	clusterCache      rocontrollers.ClusterCache
	secretCache       corecontrollers.SecretCache
	kubeconfigManager *kubeconfig.Manager
	apply             apply.Apply

	capiClusters capicontrollers.ClusterCache
	capiMachines capicontrollers.MachineCache
}

func Register(
	ctx context.Context,
	clients *wrangler.Context) {
	h := handler{
		mgmtClusterCache:  clients.Mgmt.Cluster().Cache(),
		mgmtClusters:      clients.Mgmt.Cluster(),
		clusterTokenCache: clients.Mgmt.ClusterRegistrationToken().Cache(),
		clusterTokens:     clients.Mgmt.ClusterRegistrationToken(),
		clusters:          clients.Provisioning.Cluster(),
		clusterCache:      clients.Provisioning.Cluster().Cache(),
		secretCache:       clients.Core.Secret().Cache(),
		capiClusters:      clients.CAPI.Cluster().Cache(),
		capiMachines:      clients.CAPI.Machine().Cache(),
		kubeconfigManager: kubeconfig.New(clients),
		apply: clients.Apply.WithCacheTypes(
			clients.Provisioning.Cluster(),
			clients.Mgmt.Cluster()),
	}

	mgmtcontrollers.RegisterClusterGeneratingHandler(ctx,
		clients.Mgmt.Cluster(),
		clients.Apply.WithCacheTypes(clients.Provisioning.Cluster()),
		"",
		"provisioning-cluster-create",
		h.generateProvisioning,
		nil)

	clusterCreateApply := clients.Apply.WithCacheTypes(clients.Mgmt.Cluster(),
		clients.Mgmt.ClusterRegistrationToken(),
		clients.RBAC.ClusterRoleBinding(),
		clients.Core.Namespace(),
		clients.Core.Secret())

	if features.MCM.Enabled() {
		clusterCreateApply = clusterCreateApply.WithCacheTypes(clients.Mgmt.ClusterRoleTemplateBinding())
	}

	rocontrollers.RegisterClusterGeneratingHandler(ctx,
		clients.Provisioning.Cluster(),
		clusterCreateApply,
		"Created",
		"cluster-create",
		h.generateCluster,
		&generic.GeneratingHandlerOptions{
			AllowClusterScoped: true,
		},
	)

	clients.Mgmt.Cluster().OnChange(ctx, "cluster-watch", h.createToken)
	clusterCache := clients.Provisioning.Cluster().Cache()
	relatedresource.Watch(ctx, "cluster-watch", h.clusterWatch,
		clients.Provisioning.Cluster(), clients.Mgmt.Cluster())

	clusterCache.AddIndexer(ByCluster, byClusterIndex)

	clients.Mgmt.Cluster().OnRemove(ctx, "mgmt-cluster-remove", h.OnMgmtClusterRemove)
	clients.Provisioning.Cluster().OnRemove(ctx, "provisioning-cluster-remove", h.OnClusterRemove)
}

func byClusterIndex(obj *v1.Cluster) ([]string, error) {
	if obj.Status.ClusterName == "" {
		return nil, nil
	}
	return []string{obj.Status.ClusterName}, nil
}

func (h *handler) clusterWatch(namespace, name string, obj runtime.Object) ([]relatedresource.Key, error) {
	cluster, ok := obj.(*v3.Cluster)
	if !ok {
		return nil, nil
	}
	operatorClusters, err := h.clusterCache.GetByIndex(ByCluster, cluster.Name)
	if err != nil || len(operatorClusters) == 0 {
		// ignore
		return nil, nil
	}
	return []relatedresource.Key{
		{
			Namespace: operatorClusters[0].Namespace,
			Name:      operatorClusters[0].Name,
		},
	}, nil
}

func (h *handler) isLegacyCluster(cluster interface{}) bool {
	if c, ok := cluster.(*v3.Cluster); ok {
		return mgmtNameRegexp.MatchString(c.Name)
	} else if c, ok := cluster.(*v1.Cluster); ok {
		return mgmtNameRegexp.MatchString(c.Name)
	}
	return false
}

func (h *handler) generateProvisioning(cluster *v3.Cluster, status v3.ClusterStatus) ([]runtime.Object, v3.ClusterStatus, error) {
	if !h.isLegacyCluster(cluster) || cluster.Spec.FleetWorkspaceName == "" {
		return nil, status, nil
	}
	return []runtime.Object{
		&v1.Cluster{
			ObjectMeta: metav1.ObjectMeta{
				Name:        cluster.Name,
				Namespace:   cluster.Spec.FleetWorkspaceName,
				Labels:      yaml.CleanAnnotationsForExport(cluster.Labels),
				Annotations: yaml.CleanAnnotationsForExport(cluster.Annotations),
			},
		},
	}, status, nil
}

func (h *handler) generateCluster(cluster *v1.Cluster, status v1.ClusterStatus) ([]runtime.Object, v1.ClusterStatus, error) {
	switch {
	case cluster.Spec.ClusterAPIConfig != nil:
		return h.createClusterAndDeployAgent(cluster, status)
	default:
		return h.createCluster(cluster, status, v3.ClusterSpec{
			ImportedConfig: &v3.ImportedConfig{},
		})
	}
}

func NormalizeCluster(cluster *v3.Cluster) (runtime.Object, error) {
	// We do this so that we don't clobber status because the rancher object is pretty dirty and doesn't have a status subresource
	data, err := convert.EncodeToMap(cluster)
	if err != nil {
		return nil, err
	}
	data = map[string]interface{}{
		"metadata": data["metadata"],
		"spec":     data["spec"],
	}
	data["kind"] = "Cluster"
	data["apiVersion"] = "management.cattle.io/v3"
	return &unstructured.Unstructured{Object: data}, nil
}

func (h *handler) createToken(_ string, cluster *v3.Cluster) (*v3.Cluster, error) {
	if cluster == nil {
		return cluster, nil
	}
	_, err := h.clusterTokenCache.Get(cluster.Name, "default-token")
	if apierror.IsNotFound(err) {
		_, err = h.clusterTokens.Create(&v3.ClusterRegistrationToken{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "default-token",
				Namespace: cluster.Name,
			},
			Spec: v3.ClusterRegistrationTokenSpec{
				ClusterName: cluster.Name,
			},
		})
		return cluster, err
	}
	return cluster, err
}

func (h *handler) createCluster(cluster *v1.Cluster, status v1.ClusterStatus, spec v3.ClusterSpec) ([]runtime.Object, v1.ClusterStatus, error) {
	if h.isLegacyCluster(cluster) {
		mgmtCluster, err := h.mgmtClusterCache.Get(cluster.Name)
		if err != nil {
			return nil, status, err
		}
		return h.updateStatus(nil, cluster, status, mgmtCluster)
	}
	return h.createNewCluster(cluster, status, spec)
}

func mgmtClusterName(clusterNamespace, clusterName string) string {
	hash := sha256.Sum256([]byte(clusterNamespace + "/" + clusterName))
	return name.SafeConcatName("c", "m", hex.EncodeToString(hash[:])[:8])
}

func (h *handler) createNewCluster(cluster *v1.Cluster, status v1.ClusterStatus, spec v3.ClusterSpec) ([]runtime.Object, v1.ClusterStatus, error) {
	spec.DisplayName = cluster.Name
	spec.Description = cluster.Annotations["field.cattle.io/description"]
	spec.FleetWorkspaceName = cluster.Namespace
	spec.AgentEnvVars = cluster.Spec.AgentEnvVars
	spec.DefaultPodSecurityPolicyTemplateName = cluster.Spec.DefaultPodSecurityPolicyTemplateName
	spec.DefaultClusterRoleForProjectMembers = cluster.Spec.DefaultClusterRoleForProjectMembers
	spec.EnableNetworkPolicy = cluster.Spec.EnableNetworkPolicy

	if cluster.Spec.RKEConfig != nil {
		spec.LocalClusterAuthEndpoint = v3.LocalClusterAuthEndpoint{
			FQDN:    cluster.Spec.RKEConfig.LocalClusterAuthEndpoint.FQDN,
			CACerts: cluster.Spec.RKEConfig.LocalClusterAuthEndpoint.CACerts,
			Enabled: cluster.Spec.RKEConfig.LocalClusterAuthEndpoint.Enabled,
		}
	}

	newCluster := &v3.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Name:        mgmtClusterName(cluster.Namespace, cluster.Name),
			Labels:      cluster.Labels,
			Annotations: map[string]string{},
		},
		Spec: spec,
	}

	for k, v := range cluster.Annotations {
		newCluster.Annotations[k] = v
	}

	userName, err := h.kubeconfigManager.EnsureUser(cluster.Namespace, cluster.Name)
	if err != nil {
		return nil, status, err
	}

	newCluster.Annotations[creatorIDAnn] = userName

	normalizedCluster, err := NormalizeCluster(newCluster)
	if err != nil {
		return nil, status, err
	}

	return h.updateStatus([]runtime.Object{
		normalizedCluster,
	}, cluster, status, newCluster)
}

func (h *handler) updateStatus(objs []runtime.Object, cluster *v1.Cluster, status v1.ClusterStatus, rCluster *v3.Cluster) ([]runtime.Object, v1.ClusterStatus, error) {
	ready := false
	existing, err := h.mgmtClusterCache.Get(rCluster.Name)
	if err != nil && !apierror.IsNotFound(err) {
		return nil, status, err
	} else if err == nil {
		if condition.Cond("Ready").IsTrue(existing) {
			ready = true
		}
	}

	// Never set ready back to false because we will end up deleting the secret
	status.Ready = status.Ready || ready
	status.ObservedGeneration = cluster.Generation
	status.ClusterName = rCluster.Name
	if ready {
		kstatus.SetActive(&status)
	} else {
		kstatus.SetTransitioning(&status, "")
	}

	if status.Ready {
		secret, err := h.kubeconfigManager.GetKubeConfig(cluster, status)
		if err != nil {
			return nil, status, err
		}
		if secret != nil {
			if secret.UID == "" {
				objs = append(objs, secret)
			}
			status.ClientSecretName = secret.Name

			if features.MCM.Enabled() {
				crtb, err := h.kubeconfigManager.GetCRTBForAdmin(cluster, status)
				if err != nil {
					return nil, status, err
				}
				objs = append(objs, crtb)
			} else if cluster.Namespace == "fleet-local" && cluster.Name == "local" {
				user, err := h.kubeconfigManager.EnsureUser(cluster.Namespace, cluster.Name)
				if err != nil {
					return objs, status, err
				}
				objs = append(objs, &rbacv1.ClusterRoleBinding{
					ObjectMeta: metav1.ObjectMeta{
						Name: "fleet-local-cluster-admin",
					},
					Subjects: []rbacv1.Subject{{
						Kind:     "User",
						APIGroup: rbacv1.GroupName,
						Name:     user,
					}},
					RoleRef: rbacv1.RoleRef{
						Kind: "ClusterRole",
						Name: "cluster-admin",
					},
				})
			}
		}
	}

	return objs, status, nil
}
<|endoftext|>


<commit_before>package stardb

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"time"

	"gopkg.in/mgo.v2"

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

type Farm struct {

	// Fields with  `json:"-"` are ommited for web display

	InternalID bson.ObjectId `bson:"_id,omitempty" json:"-"`
	ID         string        // Like InternalID, but in hex format. This can be exposed on the web.
	UniqueID   int           `json:"-"`
	Name       string
	Farmer     string
	Likes      int
	SaveTime   time.Time `json:"-"`
	Thumbnail  string
}

func (f *Farm) ScreenshotPath() string {
	return fmt.Sprintf("/screenshot/%v/%d.png", f.InternalID.Hex(), f.SaveTime.Unix())
}

func (f *Farm) saveGamePath() string {
	return SaveGamePath(f.InternalID.Hex(), f.SaveTime)
}

func SaveGamePath(id string, ts time.Time) string {
	return fmt.Sprintf("/saveGames/%v/%d.xml", id, ts.Unix())
}

func FarmsJSON() ([]byte, error) {
	var result []*Farm

	if err := FarmCollection.Find(nil).Limit(20).All(&result); err != nil {
		return nil, err
	}
	for _, farm := range result {
		farm.Thumbnail = farm.ScreenshotPath()
		farm.ID = farm.InternalID.Hex()
	}
	return json.Marshal(result)
}

func UpdateFarmTime(id bson.ObjectId, ts time.Time) error {
	return FarmCollection.Update(bson.M{"_id": id}, bson.M{"$set": bson.M{"savetime": ts}})
}

func FindFarm(c *mgo.Collection, uniqueIDForThisGame int, playerName, farmName string) (ret *Farm, existing bool, err error) {
	ret = &Farm{}
	q := c.Find(bson.M{
		"name":     farmName,
		"farmer":   playerName,
		"uniqueid": uniqueIDForThisGame,
	})
	if err := q.One(&ret); err != nil {
		log.Println("not found", err)

		farm := &Farm{
			Name:       farmName,
			Farmer:     playerName,
			UniqueID:   uniqueIDForThisGame,
			InternalID: bson.NewObjectId(),
			SaveTime:   time.Now(),
		}
		if err := c.Insert(farm); err != nil {
			log.Println("could not insert", err)
			return nil, false, err
		}
		log.Println("insert ok", farm.InternalID.String())
		return farm, false, nil
	}
	log.Printf("found ok %v, %v, %v", ret.Name, ret.Farmer, ret.SaveTime)

	return ret, true, nil
}

func WriteSaveFile(farm *Farm, body []byte, ts time.Time) error {
	if farm.SaveTime.IsZero() {
		return fmt.Errorf("error writing save file: unexpected zero save time")
	}

	saveFile := farm.saveGamePath()
	g, err := GFS.Create(saveFile)
	if err != nil {
		return fmt.Errorf("Error opening grid saveGames %v: %v", saveFile, err)

	}
	g.SetUploadDate(ts)
	defer g.Close()
	if _, err := g.Write(body); err != nil {
		return fmt.Errorf("Failed to write grid save file at %v: %v", saveFile, err)
	}

	log.Printf("Wrote grid saveGame file %v", saveFile)
	return nil
}

// NewScreenshotWriter saves a screenshot in GFS at screenshots/<hexid>.png
func NewScreenshotWriter(farm *Farm, ts time.Time) (io.WriteCloser, error) {
	if farm.SaveTime.IsZero() {
		return nil, fmt.Errorf("error writing screenshot: unexpected zero save time")
	}
	g, err := GFS.Create(farm.ScreenshotPath())
	if err != nil {
		return nil, err
	}
	g.SetUploadDate(ts)
	return g, nil
}
<commit_msg>Use the given timestamp for the savegame filename<commit_after>package stardb

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"time"

	"gopkg.in/mgo.v2"

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

type Farm struct {

	// Fields with  `json:"-"` are ommited for web display

	InternalID bson.ObjectId `bson:"_id,omitempty" json:"-"`
	ID         string        // Like InternalID, but in hex format. This can be exposed on the web.
	UniqueID   int           `json:"-"`
	Name       string
	Farmer     string
	Likes      int
	SaveTime   time.Time `json:"-"`
	Thumbnail  string
}

func (f *Farm) ScreenshotPath() string {
	return fmt.Sprintf("/screenshot/%v/%d.png", f.InternalID.Hex(), f.SaveTime.Unix())
}

func (f *Farm) saveGamePath() string {
	return SaveGamePath(f.InternalID.Hex(), f.SaveTime)
}

func SaveGamePath(id string, ts time.Time) string {
	return fmt.Sprintf("/saveGames/%v/%d.xml", id, ts.Unix())
}

func FarmsJSON() ([]byte, error) {
	var result []*Farm

	if err := FarmCollection.Find(nil).Limit(20).All(&result); err != nil {
		return nil, err
	}
	for _, farm := range result {
		farm.Thumbnail = farm.ScreenshotPath()
		farm.ID = farm.InternalID.Hex()
	}
	return json.Marshal(result)
}

func UpdateFarmTime(id bson.ObjectId, ts time.Time) error {
	return FarmCollection.Update(bson.M{"_id": id}, bson.M{"$set": bson.M{"savetime": ts}})
}

func FindFarm(c *mgo.Collection, uniqueIDForThisGame int, playerName, farmName string) (ret *Farm, existing bool, err error) {
	ret = &Farm{}
	q := c.Find(bson.M{
		"name":     farmName,
		"farmer":   playerName,
		"uniqueid": uniqueIDForThisGame,
	})
	if err := q.One(&ret); err != nil {
		log.Println("not found", err)

		farm := &Farm{
			Name:       farmName,
			Farmer:     playerName,
			UniqueID:   uniqueIDForThisGame,
			InternalID: bson.NewObjectId(),
			SaveTime:   time.Now(),
		}
		if err := c.Insert(farm); err != nil {
			log.Println("could not insert", err)
			return nil, false, err
		}
		log.Println("insert ok", farm.InternalID.String())
		return farm, false, nil
	}
	log.Printf("found ok %v, %v, %v", ret.Name, ret.Farmer, ret.SaveTime)

	return ret, true, nil
}

func WriteSaveFile(farm *Farm, body []byte, ts time.Time) error {
	farm.SaveTime = ts
	saveFile := farm.saveGamePath()
	g, err := GFS.Create(saveFile)
	if err != nil {
		return fmt.Errorf("Error opening grid saveGames %v: %v", saveFile, err)

	}
	g.SetUploadDate(ts)
	defer g.Close()
	if _, err := g.Write(body); err != nil {
		return fmt.Errorf("Failed to write grid save file at %v: %v", saveFile, err)
	}

	log.Printf("Wrote grid saveGame file %v", saveFile)
	return nil
}

// NewScreenshotWriter saves a screenshot in GFS at screenshots/<hexid>.png
func NewScreenshotWriter(farm *Farm, ts time.Time) (io.WriteCloser, error) {
	if farm.SaveTime.IsZero() {
		return nil, fmt.Errorf("error writing screenshot: unexpected zero save time")
	}
	g, err := GFS.Create(farm.ScreenshotPath())
	if err != nil {
		return nil, err
	}
	g.SetUploadDate(ts)
	return g, nil
}
<|endoftext|>


<commit_before>package main

import (
    "log"
    "github.com/trendrr/goshire/cheshire"
    "github.com/trendrr/goshire-shards/shards"
    "github.com/trendrr/goshire-shards/admin/balancer"
    "github.com/trendrr/goshire/cheshire/impl/gocache"
    "flag"
    // "time"
    // "fmt"
)

//command line args
var (
    configFilename = flag.String("config-filename", "config.yaml", "filename of the config")
    dataDir = flag.String("data-dir", "data", "The local directory where data should be stored")
)

func main() {
    flag.Parse()
    bootstrap := cheshire.NewBootstrapFile(*configFilename)

    //Setup our cache.  this uses the local cache
    cache := gocache.New(10, 10)
    bootstrap.AddFilters(cheshire.NewSession(cache, 3600))

    balancer.Servs.DataDir = *dataDir
    balancer.Servs.Load()

    testrt := shards.NewRouterTable("Test")
    balancer.Servs.SetRouterTable(testrt)

    //try creating entry
    entry := &shards.RouterEntry{
        Address : "localhost",
        JsonPort : 8009,
        HttpPort : 8010,
        Partitions : make([]int, 0),
    }

    log.Println("********************* ADD ENTRY")
    testrt.AddEntries(entry)

    //
    log.Println("Starting")
    // go func() {
    //     c := time.Tick(5 * time.Second)
    //     for now := range c {
    //         str := fmt.Sprintf("%v %s\n", now, "Something something")
    //         balancer.Servs.Logger.Emit("test", str)
    //         balancer.Servs.Logger.Println("TESTING LOG")
    //     }
    // }()

    //starts listening on all configured interfaces
    bootstrap.Start()

}<commit_msg>fix main<commit_after>package main

import (
    "log"
    "github.com/trendrr/goshire/cheshire"
    // "github.com/trendrr/goshire-shards/shards"
    "github.com/trendrr/goshire-shards/admin/balancer"
    "github.com/trendrr/goshire/cheshire/impl/gocache"
    "flag"
    // "time"
    // "fmt"
)

//command line args
var (
    configFilename = flag.String("config-filename", "config.yaml", "filename of the config")
    dataDir = flag.String("data-dir", "data", "The local directory where data should be stored")
)

func main() {
    flag.Parse()
    bootstrap := cheshire.NewBootstrapFile(*configFilename)

    //Setup our cache.  this uses the local cache
    cache := gocache.New(10, 10)
    bootstrap.AddFilters(cheshire.NewSession(cache, 3600))

    balancer.Servs.DataDir = *dataDir
    balancer.Servs.Load()

    // testrt := shards.NewRouterTable("Test")
    // balancer.Servs.SetRouterTable(testrt)

    // //try creating entry
    // entry := &shards.RouterEntry{
    //     Address : "localhost",
    //     JsonPort : 8009,
    //     HttpPort : 8010,
    //     Partitions : make([]int, 0),
    // }

    // log.Println("********************* ADD ENTRY")
    // testrt.AddEntries(entry)

    //
    log.Println("Starting")
    // go func() {
    //     c := time.Tick(5 * time.Second)
    //     for now := range c {
    //         str := fmt.Sprintf("%v %s\n", now, "Something something")
    //         balancer.Servs.Logger.Emit("test", str)
    //         balancer.Servs.Logger.Println("TESTING LOG")
    //     }
    // }()

    //starts listening on all configured interfaces
    bootstrap.Start()

}<|endoftext|>


<commit_before>package agent

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"net/http"
	"time"

	"github.com/jhunt/go-log"
)

func (agent *Agent) Ping() {
	if agent.Registration.URL == "" {
		log.Infof("no registration.url provided; skipping agent auto-registration")
		return
	}
	if agent.Registration.Interval <= 0 {
		log.Errorf("invalid registration.interval %d supplied; skipping agent auto-registration", agent.Registration.Interval)
		return
	}

	pool := x509.NewCertPool()
	if agent.Registration.ShieldCACert != "" {
		if ok := pool.AppendCertsFromPEM([]byte(agent.Registration.ShieldCACert)); !ok {
			log.Errorf("Invalid or malformed CA Certificate")
			return
		}
	}

	client := &http.Client{
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{
				InsecureSkipVerify: agent.Registration.SkipVerify,
				RootCAs:            pool,
			},
			Proxy:             http.ProxyFromEnvironment,
			DisableKeepAlives: true,
		},
		Timeout: 30 * time.Second,
	}

	ping := func() {
		log.Debugf("pinging shield core")
		var params = struct {
			Name string `json:"name"`
			Port int    `json:"port"`
		}{
			Name: agent.Name,
			Port: agent.Port,
		}
		b, err := json.Marshal(params)
		if err != nil {
			log.Errorf("failed to marshal /v2/agents request parameters to JSON: %s", err)
			return
		}

		log.Debugf("pre-registering with %s/v2/agents as %s", agent.Registration.URL, string(b))
		req, err := http.NewRequest("POST", agent.Registration.URL+"/v2/agents", bytes.NewBuffer(b))
		if err != nil {
			log.Errorf("failed to issue POST %s/v2/agents: %s", agent.Registration.URL, err)
			return
		}

		res, err := client.Do(req)
		if err != nil {
			log.Errorf("failed to issue POST /v2/agents: %s", err)
			return
		}

		if res.StatusCode != 200 {
			log.Errorf("pre-registration with %s failed; SHIELD Core responeded HTTP %s", agent.Registration.URL, res.Status)
			return
		}

		log.Infof("pre-registered with %s as %s (port %d)", agent.Registration.URL, agent.Name, agent.Port)
		log.Debugf("POST /v2/agents returned [%s]", res.Status)
	}

	ch := time.Tick(time.Duration(agent.Registration.Interval) * time.Second)
	ping()
	for _ = range ch {
		ping()
	}
}
<commit_msg>Use a Ticker object in agent ping()<commit_after>package agent

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"net/http"
	"time"

	"github.com/jhunt/go-log"
)

func (agent *Agent) Ping() {
	if agent.Registration.URL == "" {
		log.Infof("no registration.url provided; skipping agent auto-registration")
		return
	}
	if agent.Registration.Interval <= 0 {
		log.Errorf("invalid registration.interval %d supplied; skipping agent auto-registration", agent.Registration.Interval)
		return
	}

	pool := x509.NewCertPool()
	if agent.Registration.ShieldCACert != "" {
		if ok := pool.AppendCertsFromPEM([]byte(agent.Registration.ShieldCACert)); !ok {
			log.Errorf("Invalid or malformed CA Certificate")
			return
		}
	}

	client := &http.Client{
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{
				InsecureSkipVerify: agent.Registration.SkipVerify,
				RootCAs:            pool,
			},
			Proxy:             http.ProxyFromEnvironment,
			DisableKeepAlives: true,
		},
		Timeout: 30 * time.Second,
	}

	ping := func() {
		log.Debugf("pinging shield core")
		var params = struct {
			Name string `json:"name"`
			Port int    `json:"port"`
		}{
			Name: agent.Name,
			Port: agent.Port,
		}
		b, err := json.Marshal(params)
		if err != nil {
			log.Errorf("failed to marshal /v2/agents request parameters to JSON: %s", err)
			return
		}

		log.Debugf("pre-registering with %s/v2/agents as %s", agent.Registration.URL, string(b))
		req, err := http.NewRequest("POST", agent.Registration.URL+"/v2/agents", bytes.NewBuffer(b))
		if err != nil {
			log.Errorf("failed to issue POST %s/v2/agents: %s", agent.Registration.URL, err)
			return
		}

		res, err := client.Do(req)
		if err != nil {
			log.Errorf("failed to issue POST /v2/agents: %s", err)
			return
		}

		if res.StatusCode != 200 {
			log.Errorf("pre-registration with %s failed; SHIELD Core responeded HTTP %s", agent.Registration.URL, res.Status)
			return
		}

		log.Infof("pre-registered with %s as %s (port %d)", agent.Registration.URL, agent.Name, agent.Port)
		log.Debugf("POST /v2/agents returned [%s]", res.Status)
	}

	t := time.NewTicker(time.Duration(agent.Registration.Interval) * time.Second)
	ping()
	for _ = range ch {
		ping()
	}
}
<|endoftext|>


<commit_before>package orgnetsim

import (
	"testing"
)

type testNetwork struct {
	relatedAgents []*Agent
	agentByID     map[string]*Agent
	LinkStrength  int
}

func (tn *testNetwork) GetRelatedAgents(a *Agent) []*Agent {
	return tn.relatedAgents
}

func (tn *testNetwork) GetAgentByID(id string) *Agent {
	return tn.agentByID[id]
}

func (tn *testNetwork) IncrementLinkStrength(id1 string, id2 string) error {
	tn.LinkStrength = tn.LinkStrength + 1
	return nil
}

func (tn *testNetwork) addAgent(a *Agent) {
	tn.relatedAgents = append(tn.relatedAgents, a)
	tn.agentByID[a.ID] = a
}

func newTestNetwork() *testNetwork {
	tn := testNetwork{}
	tn.relatedAgents = []*Agent{}
	tn.agentByID = map[string]*Agent{}
	ra1 := newAgent()
	ra1.ID = "id_1"
	ra1.Influence = 0.5
	ra1.Susceptability = 0.5
	ra1.Contrariness = 0.5
	ra1.Color = Blue
	ra1.Mail <- "id_5"
	tn.addAgent(ra1)
	ra2 := newAgent()
	ra2.ID = "id_2"
	ra2.Mail <- "id_6"
	tn.addAgent(ra2)
	ra3 := newAgent()
	ra3.ID = "id_3"
	tn.addAgent(ra3)
	return &tn
}

func newAgent() *Agent {
	a := Agent{}
	a.Mail = make(chan string, 1)
	return &a
}

func TestSendMailSendsMsgToFirstAvailableRelatedAgent(t *testing.T) {
	tn := newTestNetwork()
	aut := newAgent()
	aut.ID = "id_aut"
	count := aut.SendMail(tn)
	AreEqual(t, 1, count, "Message not sent to first free Agent")
	msg, received := tn.GetAgentByID("id_3").RecieveMsg()
	IsTrue(t, received, "Message not sent to first free agent in related agents")
	AreEqual(t, "id_aut", msg, "Wrong message sent to first free agent")
}

func TestSendMailDoesNotSendIfNoAvailableRelatedAgent(t *testing.T) {
	tn := newTestNetwork()
	aut := newAgent()
	aut.ID = "id_aut"
	tn.GetAgentByID("id_3").Mail <- "block"
	count := aut.SendMail(tn)
	AreEqual(t, 0, count, "Message sent but should not have been since there are no Agents free")
}

func TestReadMailReceivesMsgIncrementsLinkStrength(t *testing.T) {
	tn := newTestNetwork()
	aut := newAgent()
	aut.ID = "id_aut"
	aut.Susceptability = 1
	aut.Color = Red
	sent := aut.SendMsg("id_1")
	AreEqual(t, 0, tn.LinkStrength, "LinkStrength not initialised to 0")
	aut.ReadMail(tn)
	IsTrue(t, sent, "Msg not sent to Agent under test")
	AreEqual(t, 1, tn.LinkStrength, "LinkStrength not incremented as part of reading a Mail")
	sent = aut.SendMsg("id_1")
	IsTrue(t, sent, "Msg not sent to Agent under test")
	aut.ReadMail(tn)
	AreEqual(t, 2, tn.LinkStrength, "LinkStrength not incremented as part of reading a Mail")
}

func TestReadMailReceivesMsgHigherSusceptabilityIgnoresMsg(t *testing.T) {
	tn := newTestNetwork()
	aut := newAgent()
	aut.ID = "id_aut"
	aut.Susceptability = 1
	aut.Color = Red
	sent := aut.SendMsg("id_1")
	aut.ReadMail(tn)
	IsTrue(t, sent, "Msg not sent to Agent under test")
	AreEqual(t, Red, aut.Color, "Agent Color should not change if Agent has higher susceptability")
}

func TestReadMailReceivesMsgLowerSusceptabilityChangesColor(t *testing.T) {
	tn := newTestNetwork()
	aut := newAgent()
	aut.ID = "id_aut"
	aut.Susceptability = 0.4
	aut.Color = Red
	sent := aut.SendMsg("id_1")
	aut.ReadMail(tn)
	IsTrue(t, sent, "Msg not sent to Agent under test")
	AreEqual(t, Blue, aut.Color, "Agent Color should change to Blue if Agent has lower susceptability")
}

func TestReadMailReceivesMsgLowerSusceptabilityHigherContrarinessRadomlyChangesColor(t *testing.T) {
	tn := newTestNetwork()
	aut := newAgent()
	aut.ID = "id_aut"
	aut.Susceptability = 0.4
	aut.Contrariness = 0.6
	aut.Color = Red
	sent := aut.SendMsg("id_1")
	aut.ReadMail(tn)
	IsTrue(t, sent, "Msg not sent to Agent under test")
	NotEqual(t, Blue, aut.Color, "Agent Color should change to random Color if Agent has higher contrariness")
	NotEqual(t, Red, aut.Color, "Agent Color should change to random Color if Agent has higher contrariness")
}

func TestRecieveMsgReturnsFalseWhenNoMsg(t *testing.T) {
	a := newAgent()
	_, received := a.RecieveMsg()
	IsFalse(t, received, "Unexpected true returned from ReceiveMsg")
}

func TestRecieveMsgGetsMsg(t *testing.T) {
	a := newAgent()
	origMsg := "myMsg"
	a.Mail <- origMsg
	msg, received := a.RecieveMsg()
	IsTrue(t, received, "Unexpected false returned from ReceiveMsg")
	AreEqual(t, origMsg, msg, "Msgs not equal.")
}

func TestSendMsgSendsMsg(t *testing.T) {
	a := newAgent()
	origMsg := "myMsg"
	sent := a.SendMsg(origMsg)
	IsTrue(t, sent, "Msg not sent")
}

func TestSendMsgFailsSecondTime(t *testing.T) {
	a := newAgent()
	origMsg := "myMsg"
	a.SendMsg(origMsg)
	sent := a.SendMsg(origMsg)
	IsFalse(t, sent, "2nd msg sent but should have returned failed and returned false")
}

func TestClearMail(t *testing.T) {
	a := newAgent()
	origMsg := "myMsg"
	a.Mail <- origMsg
	a.ClearMail()
	select {
	case msg := <-a.Mail:
		t.Errorf("%s not cleared from Mail channel", msg)
		return
	default:
		return
	}
}

func TestClearMsgDoesntBlockWhenChannelEmpty(t *testing.T) {
	a := newAgent()
	a.ClearMail()
}

func TestDefaultColorIsGrey(t *testing.T) {
	a := newAgent()
	AreEqual(t, a.Color, Grey, "Default Agent color is not Grey")
}

func TestSetColorChangesColor(t *testing.T) {
	a := newAgent()
	a.SetColor(Blue)
	AreEqual(t, a.Color, Blue, "Agent color is not set to Blue")
}

func TestSetColorIncrementsChangeCount(t *testing.T) {
	a := newAgent()
	AreEqual(t, a.ChangeCount, 0, "Default change count is not set to 0")
	a.SetColor(Blue)
	AreEqual(t, a.ChangeCount, 1, "Change count is not incremented to 1")
	a.SetColor(Red)
	AreEqual(t, a.ChangeCount, 2, "Change count is not incremented to 2")
}
<commit_msg>Fixed broken unit test<commit_after>package orgnetsim

import (
	"testing"
)

type testNetwork struct {
	relatedAgents []*Agent
	agentByID     map[string]*Agent
	LinkStrength  int
}

func (tn *testNetwork) GetRelatedAgents(a *Agent) []*Agent {
	return tn.relatedAgents
}

func (tn *testNetwork) GetAgentByID(id string) *Agent {
	return tn.agentByID[id]
}

func (tn *testNetwork) IncrementLinkStrength(id1 string, id2 string) error {
	tn.LinkStrength = tn.LinkStrength + 1
	return nil
}

func (tn *testNetwork) addAgent(a *Agent) {
	tn.relatedAgents = append(tn.relatedAgents, a)
	tn.agentByID[a.ID] = a
}

func newTestNetwork() *testNetwork {
	tn := testNetwork{}
	tn.relatedAgents = []*Agent{}
	tn.agentByID = map[string]*Agent{}
	ra1 := newAgent()
	ra1.ID = "id_1"
	ra1.Influence = 0.5
	ra1.Susceptability = 0.5
	ra1.Contrariness = 0.5
	ra1.Color = Blue
	ra1.Mail <- "id_5"
	tn.addAgent(ra1)
	ra2 := newAgent()
	ra2.ID = "id_2"
	ra2.Mail <- "id_6"
	tn.addAgent(ra2)
	ra3 := newAgent()
	ra3.ID = "id_3"
	tn.addAgent(ra3)
	return &tn
}

func newAgent() *Agent {
	a := Agent{}
	a.Mail = make(chan string, 1)
	return &a
}

func TestSendMailSendsMsgToFirstAvailableRelatedAgent(t *testing.T) {
	tn := newTestNetwork()
	aut := newAgent()
	aut.ID = "id_aut"
	count := aut.SendMail(tn)
	AreEqual(t, 1, count, "Message not sent to first free Agent")
	msg, received := tn.GetAgentByID("id_3").RecieveMsg()
	IsTrue(t, received, "Message not sent to first free agent in related agents")
	AreEqual(t, "id_aut", msg, "Wrong message sent to first free agent")
}

func TestSendMailDoesNotSendIfNoAvailableRelatedAgent(t *testing.T) {
	tn := newTestNetwork()
	aut := newAgent()
	aut.ID = "id_aut"
	tn.GetAgentByID("id_3").Mail <- "block"
	count := aut.SendMail(tn)
	AreEqual(t, 0, count, "Message sent but should not have been since there are no Agents free")
}

func TestReadMailReceivesMsgIncrementsLinkStrength(t *testing.T) {
	tn := newTestNetwork()
	aut := newAgent()
	aut.ID = "id_aut"
	aut.Susceptability = 1
	aut.Color = Red
	sent := aut.SendMsg("id_1")
	AreEqual(t, 0, tn.LinkStrength, "LinkStrength not initialised to 0")
	aut.ReadMail(tn)
	IsTrue(t, sent, "Msg not sent to Agent under test")
	AreEqual(t, 1, tn.LinkStrength, "LinkStrength not incremented as part of reading a Mail")
	sent = aut.SendMsg("id_1")
	IsTrue(t, sent, "Msg not sent to Agent under test")
	aut.ReadMail(tn)
	AreEqual(t, 2, tn.LinkStrength, "LinkStrength not incremented as part of reading a Mail")
}

func TestReadMailReceivesMsgHigherSusceptabilityIgnoresMsg(t *testing.T) {
	tn := newTestNetwork()
	aut := newAgent()
	aut.ID = "id_aut"
	aut.Susceptability = 1
	aut.Color = Red
	sent := aut.SendMsg("id_1")
	aut.ReadMail(tn)
	IsTrue(t, sent, "Msg not sent to Agent under test")
	AreEqual(t, Red, aut.Color, "Agent Color should not change if Agent has higher susceptability")
}

func TestReadMailReceivesMsgLowerSusceptabilityChangesColor(t *testing.T) {
	tn := newTestNetwork()
	aut := newAgent()
	aut.ID = "id_aut"
	aut.Susceptability = 0.4
	aut.Color = Red
	sent := aut.SendMsg("id_1")
	aut.ReadMail(tn)
	IsTrue(t, sent, "Msg not sent to Agent under test")
	AreEqual(t, Blue, aut.Color, "Agent Color should change to Blue if Agent has lower susceptability")
}

func TestReadMailReceivesMsgLowerSusceptabilityHigherContrarinessRandomlyChangesColor(t *testing.T) {
	tn := newTestNetwork()
	aut := newAgent()
	aut.ID = "id_aut"
	aut.Susceptability = 0.4
	aut.Contrariness = 0.6
	aut.Color = Red
	sent := aut.SendMsg("id_1")
	aut.ReadMail(tn)
	IsTrue(t, sent, "Msg not sent to Agent under test")
	NotEqual(t, Grey, aut.Color, "Agent Color should change to random Color other than Grey if Agent has higher contrariness")
	NotEqual(t, Red, aut.Color, "Agent Color should change to random Color other than it was if Agent has higher contrariness")
}

func TestRecieveMsgReturnsFalseWhenNoMsg(t *testing.T) {
	a := newAgent()
	_, received := a.RecieveMsg()
	IsFalse(t, received, "Unexpected true returned from ReceiveMsg")
}

func TestRecieveMsgGetsMsg(t *testing.T) {
	a := newAgent()
	origMsg := "myMsg"
	a.Mail <- origMsg
	msg, received := a.RecieveMsg()
	IsTrue(t, received, "Unexpected false returned from ReceiveMsg")
	AreEqual(t, origMsg, msg, "Msgs not equal.")
}

func TestSendMsgSendsMsg(t *testing.T) {
	a := newAgent()
	origMsg := "myMsg"
	sent := a.SendMsg(origMsg)
	IsTrue(t, sent, "Msg not sent")
}

func TestSendMsgFailsSecondTime(t *testing.T) {
	a := newAgent()
	origMsg := "myMsg"
	a.SendMsg(origMsg)
	sent := a.SendMsg(origMsg)
	IsFalse(t, sent, "2nd msg sent but should have returned failed and returned false")
}

func TestClearMail(t *testing.T) {
	a := newAgent()
	origMsg := "myMsg"
	a.Mail <- origMsg
	a.ClearMail()
	select {
	case msg := <-a.Mail:
		t.Errorf("%s not cleared from Mail channel", msg)
		return
	default:
		return
	}
}

func TestClearMsgDoesntBlockWhenChannelEmpty(t *testing.T) {
	a := newAgent()
	a.ClearMail()
}

func TestDefaultColorIsGrey(t *testing.T) {
	a := newAgent()
	AreEqual(t, a.Color, Grey, "Default Agent color is not Grey")
}

func TestSetColorChangesColor(t *testing.T) {
	a := newAgent()
	a.SetColor(Blue)
	AreEqual(t, a.Color, Blue, "Agent color is not set to Blue")
}

func TestSetColorIncrementsChangeCount(t *testing.T) {
	a := newAgent()
	AreEqual(t, a.ChangeCount, 0, "Default change count is not set to 0")
	a.SetColor(Blue)
	AreEqual(t, a.ChangeCount, 1, "Change count is not incremented to 1")
	a.SetColor(Red)
	AreEqual(t, a.ChangeCount, 2, "Change count is not incremented to 2")
}
<|endoftext|>


<commit_before>package statscollector

import (
	"strings"
	"sync"
	"time"

	"github.com/contiv/vpp/plugins/contiv"
	"github.com/golang/protobuf/proto"
	"github.com/ligato/cn-infra/datasync"
	"github.com/ligato/cn-infra/flavors/local"
	prometheusplugin "github.com/ligato/cn-infra/rpc/prometheus"
	"github.com/ligato/vpp-agent/plugins/defaultplugins/ifplugin/model/interfaces"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

const (
	// path where the statistics are exposed
	prometheusStatsPath = "/stats"

	podNameLabel       = "podName"
	podNamespaceLabel  = "podNamespace"
	interfaceNameLabel = "interfaceName"
	nodeLabel          = "node"

	inPacketsMetric       = "inPackets"
	outPacketsMetric      = "outPackets"
	inBytesMetric         = "inBytes"
	outBytesMetric        = "outBytes"
	dropPacketsMetric     = "dropPackets"
	puntPacketsMetric     = "puntPackets"
	ipv4PacketsMetric     = "ipv4Packets"
	ipv6PacketsMetric     = "ipv6Packets"
	inNobufPacketsMetric  = "inNobufPackets"
	inMissPacketsMetric   = "inMissPackets"
	inErrorPacketsMetric  = "inErrorPackets"
	outErrorPacketsMetric = "outErrorPackets"
)

// Plugin collects the statistics from vpp interfaces and publishes them to prometheus.
type Plugin struct {
	Deps
	sync.Mutex
	ifStats   map[string]*stats
	closeCh   chan interface{}
	gaugeVecs map[string]*prometheus.GaugeVec
}

type stats struct {
	podName      string
	podNamespace string
	data         *interfaces.InterfacesState_Interface
	metrics      map[string]prometheus.Gauge
}

// Deps groups the dependencies of the Plugin.
type Deps struct {
	local.PluginInfraDeps

	// Contiv plugin is used to lookup pod related to interfaces statistics
	Contiv contiv.API

	// Prometheus plugin used to stream statistics
	Prometheus prometheusplugin.API
}

// Init initializes the plugin resources
func (p *Plugin) Init() error {
	p.closeCh = make(chan interface{})
	p.ifStats = map[string]*stats{}
	p.gaugeVecs = map[string]*prometheus.GaugeVec{}

	if p.Prometheus != nil {
		// create new registry for statistics
		err := p.Prometheus.NewRegistry(prometheusStatsPath, promhttp.HandlerOpts{ErrorHandling: promhttp.ContinueOnError, ErrorLog: p.Log})
		if err != nil {
			return err
		}

		// initialize gauge vectors for statistics
		for _, statItem := range [][2]string{
			{inPacketsMetric, "Number of received packets for interface"},
			{outPacketsMetric, "Number of transmitted packets for interface"},
			{inBytesMetric, "Number of received bytes for interface"},
			{outBytesMetric, "Number of transmitted bytes for interface"},
			{dropPacketsMetric, "Number of dropped packets for interface"},
			{puntPacketsMetric, "Number of punt packets for interface"},
			{ipv4PacketsMetric, "Number of ipv4 packets for interface"},
			{ipv6PacketsMetric, "Number of ipv6 packets for interface"},
			{inNobufPacketsMetric, "Number of received packets ??? for interface"}, // TODO describe metric
			{inMissPacketsMetric, "Number of missed packets for interface"},
			{inErrorPacketsMetric, "Number of received packets with error for interface"},
			{outErrorPacketsMetric, "Number of transmitted packets with error for interface"},
		} {

			name := statItem[0]
			help := statItem[1]

			p.gaugeVecs[name] = prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Name: name,
				Help: help,
				ConstLabels: prometheus.Labels{
					nodeLabel: p.ServiceLabel.GetAgentLabel(),
				},
			}, []string{podNameLabel, podNamespaceLabel, interfaceNameLabel})

		}

		// register created vectors to prometheus
		for name, metric := range p.gaugeVecs {
			err = p.Prometheus.Register(prometheusStatsPath, metric)
			if err != nil {
				p.Log.Errorf("failed to register %v metric %v", name, err)
				return err
			}
		}

	}

	go p.PrintStats()

	return nil
}

// Close cleans up the plugin resources
func (p *Plugin) Close() error {
	close(p.closeCh)
	return nil
}

// PrintStats dumps stats to log
func (p *Plugin) PrintStats() {
	for {
		select {
		case <-p.closeCh:
			return
		case <-time.After(10 * time.Second):
			p.Lock()
			for _, v := range p.ifStats {
				p.Log.Debugf("%v %v %v %+v", v.data.Name, v.podName, v.podNamespace, *v.data.Statistics)
			}
			p.Unlock()
		}
	}
}

// Put updates the statistics for the given key
func (p *Plugin) Put(key string, data proto.Message, opts ...datasync.PutOption) error {
	p.Lock()
	defer p.Unlock()

	if strings.HasPrefix(key, interfaces.InterfaceStateKeyPrefix()) {
		var (
			err            error
			entry          *stats
			podName, podNs string
			found          bool
		)
		const contivSystemInterfacePlaceholder = "--"
		if st, ok := data.(*interfaces.InterfacesState_Interface); ok {
			entry, found = p.ifStats[key]
			// interface is associated with a pod and we're already streaming its statistics
			if found && entry.podName != "" {
				entry.data = st
				p.ifStats[key] = entry
				p.updatePrometheusStats(entry)
			} else {
				// adding stats for new interface
				contivInterface := p.isContivSystemInterface(st.Name)
				if contivInterface {
					podName = contivSystemInterfacePlaceholder
					podNs = contivSystemInterfacePlaceholder
				} else {
					podNs, podName, found = p.Contiv.GetPodByIf(st.Name)
				}

				if found || contivInterface {
					entry = &stats{
						podName:      podName,
						podNamespace: podNs,
						data:         st,
						metrics:      map[string]prometheus.Gauge{},
					}

					for k, vec := range p.gaugeVecs {
						entry.metrics[k], err = vec.GetMetricWith(prometheus.Labels{
							podNameLabel:       podName,
							podNamespaceLabel:  podNs,
							interfaceNameLabel: st.Name,
						})
						if err != nil {
							p.Log.Error(err)
						}
					}
					p.ifStats[key] = entry
					p.updatePrometheusStats(entry)
				}
			}
		} else {
			p.Log.Warn("Unable to decode received stats")
		}
	}

	return nil
}

// updatePrometheusStats publishes the statistics for the given interfaces into prometheus
func (p *Plugin) updatePrometheusStats(entry *stats) {
	if inPacket, found := entry.metrics[inPacketsMetric]; found && inPacket != nil {
		inPacket.Set(float64(entry.data.Statistics.InPackets))
	}
	if outPacket, found := entry.metrics[outPacketsMetric]; found && outPacket != nil {
		outPacket.Set(float64(entry.data.Statistics.OutPackets))
	}
	if inBytes, found := entry.metrics[inBytesMetric]; found && inBytes != nil {
		inBytes.Set(float64(entry.data.Statistics.InBytes))
	}
	if outBytes, found := entry.metrics[outBytesMetric]; found && outBytes != nil {
		outBytes.Set(float64(entry.data.Statistics.OutBytes))
	}
	if dropPacket, found := entry.metrics[dropPacketsMetric]; found && dropPacket != nil {
		dropPacket.Set(float64(entry.data.Statistics.DropPackets))
	}
	if puntPacket, found := entry.metrics[puntPacketsMetric]; found && puntPacket != nil {
		puntPacket.Set(float64(entry.data.Statistics.PuntPackets))
	}
	if ipv4Packet, found := entry.metrics[ipv4PacketsMetric]; found && ipv4Packet != nil {
		ipv4Packet.Set(float64(entry.data.Statistics.Ipv4Packets))
	}
	if ipv6Packet, found := entry.metrics[ipv6PacketsMetric]; found && ipv6Packet != nil {
		ipv6Packet.Set(float64(entry.data.Statistics.Ipv6Packets))
	}
	if inNoBufPacket, found := entry.metrics[inNobufPacketsMetric]; found && inNoBufPacket != nil {
		inNoBufPacket.Set(float64(entry.data.Statistics.InNobufPackets))
	}
	if inMissPacket, found := entry.metrics[inMissPacketsMetric]; found && inMissPacket != nil {
		inMissPacket.Set(float64(entry.data.Statistics.InMissPackets))
	}
	if inErrorPacket, found := entry.metrics[inErrorPacketsMetric]; found && inErrorPacket != nil {
		inErrorPacket.Set(float64(entry.data.Statistics.InErrorPackets))
	}
	if outErrorPacket, found := entry.metrics[outErrorPacketsMetric]; found && outErrorPacket != nil {
		outErrorPacket.Set(float64(entry.data.Statistics.OutErrorPackets))
	}
}

// isContivSystemInterface returns true if given interface name is not associated
// with a pod (e.g. interface that interconnect vpp and host stack), otherwise false
func (p *Plugin) isContivSystemInterface(ifName string) bool {
	for _, n := range []string{"afpacket-vpp2", "vpp2", "vxlanBVI", "loopbackNIC"} {
		if n == ifName {
			return true
		}
	}
	return false
}
<commit_msg>Todo added<commit_after>package statscollector

import (
	"strings"
	"sync"
	"time"

	"github.com/contiv/vpp/plugins/contiv"
	"github.com/golang/protobuf/proto"
	"github.com/ligato/cn-infra/datasync"
	"github.com/ligato/cn-infra/flavors/local"
	prometheusplugin "github.com/ligato/cn-infra/rpc/prometheus"
	"github.com/ligato/vpp-agent/plugins/defaultplugins/ifplugin/model/interfaces"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

const (
	// path where the statistics are exposed
	prometheusStatsPath = "/stats"

	podNameLabel       = "podName"
	podNamespaceLabel  = "podNamespace"
	interfaceNameLabel = "interfaceName"
	nodeLabel          = "node"

	inPacketsMetric       = "inPackets"
	outPacketsMetric      = "outPackets"
	inBytesMetric         = "inBytes"
	outBytesMetric        = "outBytes"
	dropPacketsMetric     = "dropPackets"
	puntPacketsMetric     = "puntPackets"
	ipv4PacketsMetric     = "ipv4Packets"
	ipv6PacketsMetric     = "ipv6Packets"
	inNobufPacketsMetric  = "inNobufPackets"
	inMissPacketsMetric   = "inMissPackets"
	inErrorPacketsMetric  = "inErrorPackets"
	outErrorPacketsMetric = "outErrorPackets"
)

// Plugin collects the statistics from vpp interfaces and publishes them to prometheus.
type Plugin struct {
	Deps
	sync.Mutex
	ifStats   map[string]*stats
	closeCh   chan interface{}
	gaugeVecs map[string]*prometheus.GaugeVec
}

type stats struct {
	podName      string
	podNamespace string
	data         *interfaces.InterfacesState_Interface
	metrics      map[string]prometheus.Gauge
}

// Deps groups the dependencies of the Plugin.
type Deps struct {
	local.PluginInfraDeps

	// Contiv plugin is used to lookup pod related to interfaces statistics
	Contiv contiv.API

	// Prometheus plugin used to stream statistics
	Prometheus prometheusplugin.API
}

// Init initializes the plugin resources
func (p *Plugin) Init() error {
	p.closeCh = make(chan interface{})
	p.ifStats = map[string]*stats{}
	p.gaugeVecs = map[string]*prometheus.GaugeVec{}

	if p.Prometheus != nil {
		// create new registry for statistics
		err := p.Prometheus.NewRegistry(prometheusStatsPath, promhttp.HandlerOpts{ErrorHandling: promhttp.ContinueOnError, ErrorLog: p.Log})
		if err != nil {
			return err
		}

		// initialize gauge vectors for statistics
		for _, statItem := range [][2]string{
			{inPacketsMetric, "Number of received packets for interface"},
			{outPacketsMetric, "Number of transmitted packets for interface"},
			{inBytesMetric, "Number of received bytes for interface"},
			{outBytesMetric, "Number of transmitted bytes for interface"},
			{dropPacketsMetric, "Number of dropped packets for interface"},
			{puntPacketsMetric, "Number of punt packets for interface"},
			{ipv4PacketsMetric, "Number of ipv4 packets for interface"},
			{ipv6PacketsMetric, "Number of ipv6 packets for interface"},
			{inNobufPacketsMetric, "Number of received packets ??? for interface"}, // TODO describe metric
			{inMissPacketsMetric, "Number of missed packets for interface"},
			{inErrorPacketsMetric, "Number of received packets with error for interface"},
			{outErrorPacketsMetric, "Number of transmitted packets with error for interface"},
		} {

			name := statItem[0]
			help := statItem[1]

			p.gaugeVecs[name] = prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Name: name,
				Help: help,
				ConstLabels: prometheus.Labels{
					nodeLabel: p.ServiceLabel.GetAgentLabel(),
				},
			}, []string{podNameLabel, podNamespaceLabel, interfaceNameLabel})

		}

		// register created vectors to prometheus
		for name, metric := range p.gaugeVecs {
			err = p.Prometheus.Register(prometheusStatsPath, metric)
			if err != nil {
				p.Log.Errorf("failed to register %v metric %v", name, err)
				return err
			}
		}

	}

	// TODO watch containerIDX and remove gauges of pods that have been deleted

	go p.PrintStats()

	return nil
}

// Close cleans up the plugin resources
func (p *Plugin) Close() error {
	close(p.closeCh)
	return nil
}

// PrintStats dumps stats to log
func (p *Plugin) PrintStats() {
	for {
		select {
		case <-p.closeCh:
			return
		case <-time.After(10 * time.Second):
			p.Lock()
			for _, v := range p.ifStats {
				p.Log.Debugf("%v %v %v %+v", v.data.Name, v.podName, v.podNamespace, *v.data.Statistics)
			}
			p.Unlock()
		}
	}
}

// Put updates the statistics for the given key
func (p *Plugin) Put(key string, data proto.Message, opts ...datasync.PutOption) error {
	p.Lock()
	defer p.Unlock()

	if strings.HasPrefix(key, interfaces.InterfaceStateKeyPrefix()) {
		var (
			err            error
			entry          *stats
			podName, podNs string
			found          bool
		)
		const contivSystemInterfacePlaceholder = "--"
		if st, ok := data.(*interfaces.InterfacesState_Interface); ok {
			entry, found = p.ifStats[key]
			// interface is associated with a pod and we're already streaming its statistics
			if found && entry.podName != "" {
				entry.data = st
				p.ifStats[key] = entry
				p.updatePrometheusStats(entry)
			} else {
				// adding stats for new interface
				contivInterface := p.isContivSystemInterface(st.Name)
				if contivInterface {
					podName = contivSystemInterfacePlaceholder
					podNs = contivSystemInterfacePlaceholder
				} else {
					podNs, podName, found = p.Contiv.GetPodByIf(st.Name)
				}

				if found || contivInterface {
					entry = &stats{
						podName:      podName,
						podNamespace: podNs,
						data:         st,
						metrics:      map[string]prometheus.Gauge{},
					}

					for k, vec := range p.gaugeVecs {
						entry.metrics[k], err = vec.GetMetricWith(prometheus.Labels{
							podNameLabel:       podName,
							podNamespaceLabel:  podNs,
							interfaceNameLabel: st.Name,
						})
						if err != nil {
							p.Log.Error(err)
						}
					}
					p.ifStats[key] = entry
					p.updatePrometheusStats(entry)
				}
			}
		} else {
			p.Log.Warn("Unable to decode received stats")
		}
	}

	return nil
}

// updatePrometheusStats publishes the statistics for the given interfaces into prometheus
func (p *Plugin) updatePrometheusStats(entry *stats) {
	if inPacket, found := entry.metrics[inPacketsMetric]; found && inPacket != nil {
		inPacket.Set(float64(entry.data.Statistics.InPackets))
	}
	if outPacket, found := entry.metrics[outPacketsMetric]; found && outPacket != nil {
		outPacket.Set(float64(entry.data.Statistics.OutPackets))
	}
	if inBytes, found := entry.metrics[inBytesMetric]; found && inBytes != nil {
		inBytes.Set(float64(entry.data.Statistics.InBytes))
	}
	if outBytes, found := entry.metrics[outBytesMetric]; found && outBytes != nil {
		outBytes.Set(float64(entry.data.Statistics.OutBytes))
	}
	if dropPacket, found := entry.metrics[dropPacketsMetric]; found && dropPacket != nil {
		dropPacket.Set(float64(entry.data.Statistics.DropPackets))
	}
	if puntPacket, found := entry.metrics[puntPacketsMetric]; found && puntPacket != nil {
		puntPacket.Set(float64(entry.data.Statistics.PuntPackets))
	}
	if ipv4Packet, found := entry.metrics[ipv4PacketsMetric]; found && ipv4Packet != nil {
		ipv4Packet.Set(float64(entry.data.Statistics.Ipv4Packets))
	}
	if ipv6Packet, found := entry.metrics[ipv6PacketsMetric]; found && ipv6Packet != nil {
		ipv6Packet.Set(float64(entry.data.Statistics.Ipv6Packets))
	}
	if inNoBufPacket, found := entry.metrics[inNobufPacketsMetric]; found && inNoBufPacket != nil {
		inNoBufPacket.Set(float64(entry.data.Statistics.InNobufPackets))
	}
	if inMissPacket, found := entry.metrics[inMissPacketsMetric]; found && inMissPacket != nil {
		inMissPacket.Set(float64(entry.data.Statistics.InMissPackets))
	}
	if inErrorPacket, found := entry.metrics[inErrorPacketsMetric]; found && inErrorPacket != nil {
		inErrorPacket.Set(float64(entry.data.Statistics.InErrorPackets))
	}
	if outErrorPacket, found := entry.metrics[outErrorPacketsMetric]; found && outErrorPacket != nil {
		outErrorPacket.Set(float64(entry.data.Statistics.OutErrorPackets))
	}
}

// isContivSystemInterface returns true if given interface name is not associated
// with a pod (e.g. interface that interconnect vpp and host stack), otherwise false
func (p *Plugin) isContivSystemInterface(ifName string) bool {
	for _, n := range []string{"afpacket-vpp2", "vpp2", "vxlanBVI", "loopbackNIC"} {
		if n == ifName {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package mocks

import (
	"sync"
	"sync/atomic"

	"github.com/Shopify/sarama"
)

// Consumer implements sarama's Consumer interface for testing purposes.
// Before you can start consuming from this consumer, you have to register
// topic/partitions using ExpectConsumePartition, and set expectations on them.
type Consumer struct {
	l                  sync.Mutex
	t                  ErrorReporter
	config             *sarama.Config
	partitionConsumers map[string]map[int32]*PartitionConsumer
	metadata           map[string][]int32
}

// NewConsumer returns a new mock Consumer instance. The t argument should
// be the *testing.T instance of your test method. An error will be written to it if
// an expectation is violated. The config argument is currently unused and can be set to nil.
func NewConsumer(t ErrorReporter, config *sarama.Config) *Consumer {
	if config == nil {
		config = sarama.NewConfig()
	}

	c := &Consumer{
		t:                  t,
		config:             config,
		partitionConsumers: make(map[string]map[int32]*PartitionConsumer),
	}
	return c
}

///////////////////////////////////////////////////
// Consumer interface implementation
///////////////////////////////////////////////////

// ConsumePartition implements the ConsumePartition method from the sarama.Consumer interface.
// Before you can start consuming a partition, you have to set expectations on it using
// ExpectConsumePartition. You can only consume a partition once per consumer.
func (c *Consumer) ConsumePartition(topic string, partition int32, offset int64) (sarama.PartitionConsumer, error) {
	c.l.Lock()
	defer c.l.Unlock()

	if c.partitionConsumers[topic] == nil || c.partitionConsumers[topic][partition] == nil {
		c.t.Errorf("No expectations set for %s/%d", topic, partition)
		return nil, errOutOfExpectations
	}

	pc := c.partitionConsumers[topic][partition]
	if pc.consumed {
		return nil, sarama.ConfigurationError("The topic/partition is already being consumed")
	}

	if pc.offset != AnyOffset && pc.offset != offset {
		c.t.Errorf("Unexpected offset when calling ConsumePartition for %s/%d. Expected %d, got %d.", topic, partition, pc.offset, offset)
	}

	pc.consumed = true
	return pc, nil
}

// Topics returns a list of topics, as registered with SetMetadata
func (c *Consumer) Topics() ([]string, error) {
	c.l.Lock()
	defer c.l.Unlock()

	if c.metadata == nil {
		c.t.Errorf("Unexpected call to Topics. Initialize the mock's topic metadata with SetMetadata.")
		return nil, sarama.ErrOutOfBrokers
	}

	var result []string
	for topic, _ := range c.metadata {
		result = append(result, topic)
	}
	return result, nil
}

// Partitions returns the list of parititons for the given topic, as registered with SetMetadata
func (c *Consumer) Partitions(topic string) ([]int32, error) {
	c.l.Lock()
	defer c.l.Unlock()

	if c.metadata == nil {
		c.t.Errorf("Unexpected call to Partitions. Initialize the mock's topic metadata with SetMetadata.")
		return nil, sarama.ErrOutOfBrokers
	}
	if c.metadata[topic] == nil {
		return nil, sarama.ErrUnknownTopicOrPartition
	}

	return c.metadata[topic], nil
}

// Close implements the Close method from the sarama.Consumer interface. It will close
// all registered PartitionConsumer instances.
func (c *Consumer) Close() error {
	c.l.Lock()
	defer c.l.Unlock()

	for _, partitions := range c.partitionConsumers {
		for _, partitionConsumer := range partitions {
			_ = partitionConsumer.Close()
		}
	}

	return nil
}

///////////////////////////////////////////////////
// Expectation API
///////////////////////////////////////////////////

// SetMetadata sets the clusters topic/partition metadata,
// which will be returned by Topics() and Partitions().
func (c *Consumer) SetTopicMetadata(metadata map[string][]int32) {
	c.l.Lock()
	defer c.l.Unlock()

	c.metadata = metadata
}

// ExpectConsumePartition will register a topic/partition, so you can set expectations on it.
// The registered PartitionConsumer will be returned, so you can set expectations
// on it using method chanining. Once a topic/partition is registered, you are
// expected to start consuming it using ConsumePartition. If that doesn't happen,
// an error will be written to the error reporter once the mock consumer is closed. It will
// also expect that the
func (c *Consumer) ExpectConsumePartition(topic string, partition int32, offset int64) *PartitionConsumer {
	c.l.Lock()
	defer c.l.Unlock()

	if c.partitionConsumers[topic] == nil {
		c.partitionConsumers[topic] = make(map[int32]*PartitionConsumer)
	}

	if c.partitionConsumers[topic][partition] == nil {
		c.partitionConsumers[topic][partition] = &PartitionConsumer{
			t:         c.t,
			topic:     topic,
			partition: partition,
			offset:    offset,
			messages:  make(chan *sarama.ConsumerMessage, c.config.ChannelBufferSize),
			errors:    make(chan *sarama.ConsumerError, c.config.ChannelBufferSize),
		}
	}

	return c.partitionConsumers[topic][partition]
}

///////////////////////////////////////////////////
// PartitionConsumer mock type
///////////////////////////////////////////////////

// PartitionConsumer implements sarama's PartitionConsumer interface for testing purposes.
// It is returned by the mock Consumers ConsumePartitionMethod, but only if it is
// registered first using the Consumer's ExpectConsumePartition method. Before consuming the
// Errors and Messages channel, you should specify what values will be provided on these
// channels using YieldMessage and YieldError.
type PartitionConsumer struct {
	l                       sync.Mutex
	t                       ErrorReporter
	topic                   string
	partition               int32
	offset                  int64
	messages                chan *sarama.ConsumerMessage
	errors                  chan *sarama.ConsumerError
	singleClose             sync.Once
	consumed                bool
	errorsShouldBeDrained   bool
	messagesShouldBeDrained bool
	highWaterMarkOffset     int64
}

///////////////////////////////////////////////////
// PartitionConsumer interface implementation
///////////////////////////////////////////////////

// AsyncClose implements the AsyncClose method from the sarama.PartitionConsumer interface.
func (pc *PartitionConsumer) AsyncClose() {
	pc.singleClose.Do(func() {
		close(pc.messages)
		close(pc.errors)
	})
}

// Close implements the Close method from the sarama.PartitionConsumer interface. It will
// verify whether the partition consumer was actually started.
func (pc *PartitionConsumer) Close() error {
	if !pc.consumed {
		pc.t.Errorf("Expectations set on %s/%d, but no partition consumer was started.", pc.topic, pc.partition)
		return errPartitionConsumerNotStarted
	}

	if pc.errorsShouldBeDrained && len(pc.errors) > 0 {
		pc.t.Errorf("Expected the errors channel for %s/%d to be drained on close, but found %d errors.", pc.topic, pc.partition, len(pc.errors))
	}

	if pc.messagesShouldBeDrained && len(pc.messages) > 0 {
		pc.t.Errorf("Expected the messages channel for %s/%d to be drained on close, but found %d messages.", pc.topic, pc.partition, len(pc.messages))
	}

	pc.AsyncClose()

	var (
		closeErr error
		wg       sync.WaitGroup
	)

	wg.Add(1)
	go func() {
		defer wg.Done()

		var errs = make(sarama.ConsumerErrors, 0)
		for err := range pc.errors {
			errs = append(errs, err)
		}

		if len(errs) > 0 {
			closeErr = errs
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		for _ = range pc.messages {
			// drain
		}
	}()

	wg.Wait()
	return closeErr
}

// Errors implements the Errors method from the sarama.PartitionConsumer interface.
func (pc *PartitionConsumer) Errors() <-chan *sarama.ConsumerError {
	return pc.errors
}

// Messages implements the Messages method from the sarama.PartitionConsumer interface.
func (pc *PartitionConsumer) Messages() <-chan *sarama.ConsumerMessage {
	return pc.messages
}

func (pc *PartitionConsumer) HighWaterMarkOffset() int64 {
	return atomic.LoadInt64(&pc.highWaterMarkOffset) + 1
}

///////////////////////////////////////////////////
// Expectation API
///////////////////////////////////////////////////

// YieldMessage will yield a messages Messages channel of this partition consumer
// when it is consumed. By default, the mock consumer will not verify whether this
// message was consumed from the Messages channel, because there are legitimate
// reasons forthis not to happen. ou can call ExpectMessagesDrainedOnClose so it will
// verify that the channel is empty on close.
func (pc *PartitionConsumer) YieldMessage(msg *sarama.ConsumerMessage) {
	pc.l.Lock()
	defer pc.l.Unlock()

	atomic.AddInt64(&pc.highWaterMarkOffset, 1)
	msg.Topic = pc.topic
	msg.Partition = pc.partition
	msg.Offset = atomic.LoadInt64(&pc.highWaterMarkOffset)

	pc.messages <- msg
}

// YieldError will yield an error on the Errors channel of this partition consumer
// when it is consumed. By default, the mock consumer will not verify whether this error was
// consumed from the Errors channel, because there are legitimate reasons for this
// not to happen. You can call ExpectErrorsDrainedOnClose so it will verify that
// the channel is empty on close.
func (pc *PartitionConsumer) YieldError(err error) {
	pc.errors <- &sarama.ConsumerError{
		Topic:     pc.topic,
		Partition: pc.partition,
		Err:       err,
	}
}

// ExpectMessagesDrainedOnClose sets an expectation on the partition consumer
// that the messages channel will be fully drained when Close is called. If this
// expectation is not met, an error is reported to the error reporter.
func (pc *PartitionConsumer) ExpectMessagesDrainedOnClose() {
	pc.messagesShouldBeDrained = true
}

// ExpectErrorsDrainedOnClose sets an expectation on the partition consumer
// that the errors channel will be fully drained when Close is called. If this
// expectation is not met, an error is reported to the error reporter.
func (pc *PartitionConsumer) ExpectErrorsDrainedOnClose() {
	pc.errorsShouldBeDrained = true
}
<commit_msg>Collapse two sync.atomic calls<commit_after>package mocks

import (
	"sync"
	"sync/atomic"

	"github.com/Shopify/sarama"
)

// Consumer implements sarama's Consumer interface for testing purposes.
// Before you can start consuming from this consumer, you have to register
// topic/partitions using ExpectConsumePartition, and set expectations on them.
type Consumer struct {
	l                  sync.Mutex
	t                  ErrorReporter
	config             *sarama.Config
	partitionConsumers map[string]map[int32]*PartitionConsumer
	metadata           map[string][]int32
}

// NewConsumer returns a new mock Consumer instance. The t argument should
// be the *testing.T instance of your test method. An error will be written to it if
// an expectation is violated. The config argument is currently unused and can be set to nil.
func NewConsumer(t ErrorReporter, config *sarama.Config) *Consumer {
	if config == nil {
		config = sarama.NewConfig()
	}

	c := &Consumer{
		t:                  t,
		config:             config,
		partitionConsumers: make(map[string]map[int32]*PartitionConsumer),
	}
	return c
}

///////////////////////////////////////////////////
// Consumer interface implementation
///////////////////////////////////////////////////

// ConsumePartition implements the ConsumePartition method from the sarama.Consumer interface.
// Before you can start consuming a partition, you have to set expectations on it using
// ExpectConsumePartition. You can only consume a partition once per consumer.
func (c *Consumer) ConsumePartition(topic string, partition int32, offset int64) (sarama.PartitionConsumer, error) {
	c.l.Lock()
	defer c.l.Unlock()

	if c.partitionConsumers[topic] == nil || c.partitionConsumers[topic][partition] == nil {
		c.t.Errorf("No expectations set for %s/%d", topic, partition)
		return nil, errOutOfExpectations
	}

	pc := c.partitionConsumers[topic][partition]
	if pc.consumed {
		return nil, sarama.ConfigurationError("The topic/partition is already being consumed")
	}

	if pc.offset != AnyOffset && pc.offset != offset {
		c.t.Errorf("Unexpected offset when calling ConsumePartition for %s/%d. Expected %d, got %d.", topic, partition, pc.offset, offset)
	}

	pc.consumed = true
	return pc, nil
}

// Topics returns a list of topics, as registered with SetMetadata
func (c *Consumer) Topics() ([]string, error) {
	c.l.Lock()
	defer c.l.Unlock()

	if c.metadata == nil {
		c.t.Errorf("Unexpected call to Topics. Initialize the mock's topic metadata with SetMetadata.")
		return nil, sarama.ErrOutOfBrokers
	}

	var result []string
	for topic, _ := range c.metadata {
		result = append(result, topic)
	}
	return result, nil
}

// Partitions returns the list of parititons for the given topic, as registered with SetMetadata
func (c *Consumer) Partitions(topic string) ([]int32, error) {
	c.l.Lock()
	defer c.l.Unlock()

	if c.metadata == nil {
		c.t.Errorf("Unexpected call to Partitions. Initialize the mock's topic metadata with SetMetadata.")
		return nil, sarama.ErrOutOfBrokers
	}
	if c.metadata[topic] == nil {
		return nil, sarama.ErrUnknownTopicOrPartition
	}

	return c.metadata[topic], nil
}

// Close implements the Close method from the sarama.Consumer interface. It will close
// all registered PartitionConsumer instances.
func (c *Consumer) Close() error {
	c.l.Lock()
	defer c.l.Unlock()

	for _, partitions := range c.partitionConsumers {
		for _, partitionConsumer := range partitions {
			_ = partitionConsumer.Close()
		}
	}

	return nil
}

///////////////////////////////////////////////////
// Expectation API
///////////////////////////////////////////////////

// SetMetadata sets the clusters topic/partition metadata,
// which will be returned by Topics() and Partitions().
func (c *Consumer) SetTopicMetadata(metadata map[string][]int32) {
	c.l.Lock()
	defer c.l.Unlock()

	c.metadata = metadata
}

// ExpectConsumePartition will register a topic/partition, so you can set expectations on it.
// The registered PartitionConsumer will be returned, so you can set expectations
// on it using method chanining. Once a topic/partition is registered, you are
// expected to start consuming it using ConsumePartition. If that doesn't happen,
// an error will be written to the error reporter once the mock consumer is closed. It will
// also expect that the
func (c *Consumer) ExpectConsumePartition(topic string, partition int32, offset int64) *PartitionConsumer {
	c.l.Lock()
	defer c.l.Unlock()

	if c.partitionConsumers[topic] == nil {
		c.partitionConsumers[topic] = make(map[int32]*PartitionConsumer)
	}

	if c.partitionConsumers[topic][partition] == nil {
		c.partitionConsumers[topic][partition] = &PartitionConsumer{
			t:         c.t,
			topic:     topic,
			partition: partition,
			offset:    offset,
			messages:  make(chan *sarama.ConsumerMessage, c.config.ChannelBufferSize),
			errors:    make(chan *sarama.ConsumerError, c.config.ChannelBufferSize),
		}
	}

	return c.partitionConsumers[topic][partition]
}

///////////////////////////////////////////////////
// PartitionConsumer mock type
///////////////////////////////////////////////////

// PartitionConsumer implements sarama's PartitionConsumer interface for testing purposes.
// It is returned by the mock Consumers ConsumePartitionMethod, but only if it is
// registered first using the Consumer's ExpectConsumePartition method. Before consuming the
// Errors and Messages channel, you should specify what values will be provided on these
// channels using YieldMessage and YieldError.
type PartitionConsumer struct {
	l                       sync.Mutex
	t                       ErrorReporter
	topic                   string
	partition               int32
	offset                  int64
	messages                chan *sarama.ConsumerMessage
	errors                  chan *sarama.ConsumerError
	singleClose             sync.Once
	consumed                bool
	errorsShouldBeDrained   bool
	messagesShouldBeDrained bool
	highWaterMarkOffset     int64
}

///////////////////////////////////////////////////
// PartitionConsumer interface implementation
///////////////////////////////////////////////////

// AsyncClose implements the AsyncClose method from the sarama.PartitionConsumer interface.
func (pc *PartitionConsumer) AsyncClose() {
	pc.singleClose.Do(func() {
		close(pc.messages)
		close(pc.errors)
	})
}

// Close implements the Close method from the sarama.PartitionConsumer interface. It will
// verify whether the partition consumer was actually started.
func (pc *PartitionConsumer) Close() error {
	if !pc.consumed {
		pc.t.Errorf("Expectations set on %s/%d, but no partition consumer was started.", pc.topic, pc.partition)
		return errPartitionConsumerNotStarted
	}

	if pc.errorsShouldBeDrained && len(pc.errors) > 0 {
		pc.t.Errorf("Expected the errors channel for %s/%d to be drained on close, but found %d errors.", pc.topic, pc.partition, len(pc.errors))
	}

	if pc.messagesShouldBeDrained && len(pc.messages) > 0 {
		pc.t.Errorf("Expected the messages channel for %s/%d to be drained on close, but found %d messages.", pc.topic, pc.partition, len(pc.messages))
	}

	pc.AsyncClose()

	var (
		closeErr error
		wg       sync.WaitGroup
	)

	wg.Add(1)
	go func() {
		defer wg.Done()

		var errs = make(sarama.ConsumerErrors, 0)
		for err := range pc.errors {
			errs = append(errs, err)
		}

		if len(errs) > 0 {
			closeErr = errs
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		for _ = range pc.messages {
			// drain
		}
	}()

	wg.Wait()
	return closeErr
}

// Errors implements the Errors method from the sarama.PartitionConsumer interface.
func (pc *PartitionConsumer) Errors() <-chan *sarama.ConsumerError {
	return pc.errors
}

// Messages implements the Messages method from the sarama.PartitionConsumer interface.
func (pc *PartitionConsumer) Messages() <-chan *sarama.ConsumerMessage {
	return pc.messages
}

func (pc *PartitionConsumer) HighWaterMarkOffset() int64 {
	return atomic.LoadInt64(&pc.highWaterMarkOffset) + 1
}

///////////////////////////////////////////////////
// Expectation API
///////////////////////////////////////////////////

// YieldMessage will yield a messages Messages channel of this partition consumer
// when it is consumed. By default, the mock consumer will not verify whether this
// message was consumed from the Messages channel, because there are legitimate
// reasons forthis not to happen. ou can call ExpectMessagesDrainedOnClose so it will
// verify that the channel is empty on close.
func (pc *PartitionConsumer) YieldMessage(msg *sarama.ConsumerMessage) {
	pc.l.Lock()
	defer pc.l.Unlock()

	msg.Topic = pc.topic
	msg.Partition = pc.partition
	msg.Offset = atomic.AddInt64(&pc.highWaterMarkOffset, 1)

	pc.messages <- msg
}

// YieldError will yield an error on the Errors channel of this partition consumer
// when it is consumed. By default, the mock consumer will not verify whether this error was
// consumed from the Errors channel, because there are legitimate reasons for this
// not to happen. You can call ExpectErrorsDrainedOnClose so it will verify that
// the channel is empty on close.
func (pc *PartitionConsumer) YieldError(err error) {
	pc.errors <- &sarama.ConsumerError{
		Topic:     pc.topic,
		Partition: pc.partition,
		Err:       err,
	}
}

// ExpectMessagesDrainedOnClose sets an expectation on the partition consumer
// that the messages channel will be fully drained when Close is called. If this
// expectation is not met, an error is reported to the error reporter.
func (pc *PartitionConsumer) ExpectMessagesDrainedOnClose() {
	pc.messagesShouldBeDrained = true
}

// ExpectErrorsDrainedOnClose sets an expectation on the partition consumer
// that the errors channel will be fully drained when Close is called. If this
// expectation is not met, an error is reported to the error reporter.
func (pc *PartitionConsumer) ExpectErrorsDrainedOnClose() {
	pc.errorsShouldBeDrained = true
}
<|endoftext|>