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

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

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

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

package predicates

import (
	"fmt"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/plugin/pkg/scheduler/algorithm"
)

type NodeInfo interface {
	GetNodeInfo(nodeID string) (*api.Node, error)
}

type StaticNodeInfo struct {
	*api.NodeList
}

func (nodes StaticNodeInfo) GetNodeInfo(nodeID string) (*api.Node, error) {
	for ix := range nodes.Items {
		if nodes.Items[ix].Name == nodeID {
			return &nodes.Items[ix], nil
		}
	}
	return nil, fmt.Errorf("failed to find node: %s, %#v", nodeID, nodes)
}

type ClientNodeInfo struct {
	*client.Client
}

func (nodes ClientNodeInfo) GetNodeInfo(nodeID string) (*api.Node, error) {
	return nodes.Nodes().Get(nodeID)
}

func isVolumeConflict(volume api.Volume, pod *api.Pod) bool {
	if volume.GCEPersistentDisk != nil {
		disk := volume.GCEPersistentDisk

		manifest := &(pod.Spec)
		for ix := range manifest.Volumes {
			if manifest.Volumes[ix].GCEPersistentDisk != nil &&
				manifest.Volumes[ix].GCEPersistentDisk.PDName == disk.PDName &&
				!(manifest.Volumes[ix].GCEPersistentDisk.ReadOnly && disk.ReadOnly) {
				return true
			}
		}
	}
	if volume.AWSElasticBlockStore != nil {
		volumeID := volume.AWSElasticBlockStore.VolumeID

		manifest := &(pod.Spec)
		for ix := range manifest.Volumes {
			if manifest.Volumes[ix].AWSElasticBlockStore != nil &&
				manifest.Volumes[ix].AWSElasticBlockStore.VolumeID == volumeID {
				return true
			}
		}
	}
	return false
}

// NoDiskConflict evaluates if a pod can fit due to the volumes it requests, and those that
// are already mounted. Some times of volumes are mounted onto node machines.  For now, these mounts
// are exclusive so if there is already a volume mounted on that node, another pod can't schedule
// there. This is GCE specific for now.
// TODO: migrate this into some per-volume specific code?
func NoDiskConflict(pod *api.Pod, existingPods []*api.Pod, node string) (bool, error) {
	manifest := &(pod.Spec)
	for ix := range manifest.Volumes {
		for podIx := range existingPods {
			if isVolumeConflict(manifest.Volumes[ix], existingPods[podIx]) {
				return false, nil
			}
		}
	}
	return true, nil
}

type ResourceFit struct {
	info NodeInfo
}

type resourceRequest struct {
	milliCPU int64
	memory   int64
}

func getResourceRequest(pod *api.Pod) resourceRequest {
	result := resourceRequest{}
	for ix := range pod.Spec.Containers {
		limits := pod.Spec.Containers[ix].Resources.Limits
		result.memory += limits.Memory().Value()
		result.milliCPU += limits.Cpu().MilliValue()
	}
	return result
}

func CheckPodsExceedingCapacity(pods []*api.Pod, capacity api.ResourceList) (fitting []*api.Pod, notFitting []*api.Pod) {
	totalMilliCPU := capacity.Cpu().MilliValue()
	totalMemory := capacity.Memory().Value()
	milliCPURequested := int64(0)
	memoryRequested := int64(0)
	for _, pod := range pods {
		podRequest := getResourceRequest(pod)
		fitsCPU := totalMilliCPU == 0 || (totalMilliCPU-milliCPURequested) >= podRequest.milliCPU
		fitsMemory := totalMemory == 0 || (totalMemory-memoryRequested) >= podRequest.memory
		if !fitsCPU || !fitsMemory {
			// the pod doesn't fit
			notFitting = append(notFitting, pod)
			continue
		}
		// the pod fits
		milliCPURequested += podRequest.milliCPU
		memoryRequested += podRequest.memory
		fitting = append(fitting, pod)
	}
	return
}

// PodFitsResources calculates fit based on requested, rather than used resources
func (r *ResourceFit) PodFitsResources(pod *api.Pod, existingPods []*api.Pod, node string) (bool, error) {
	podRequest := getResourceRequest(pod)
	info, err := r.info.GetNodeInfo(node)
	if err != nil {
		return false, err
	}
	if podRequest.milliCPU == 0 && podRequest.memory == 0 {
		return int64(len(existingPods)) < info.Status.Capacity.Pods().Value(), nil
	}
	pods := []*api.Pod{}
	copy(pods, existingPods)
	pods = append(existingPods, pod)
	_, exceeding := CheckPodsExceedingCapacity(pods, info.Status.Capacity)
	if len(exceeding) > 0 || int64(len(pods)) > info.Status.Capacity.Pods().Value() {
		return false, nil
	}
	return true, nil
}

func NewResourceFitPredicate(info NodeInfo) algorithm.FitPredicate {
	fit := &ResourceFit{
		info: info,
	}
	return fit.PodFitsResources
}

func NewSelectorMatchPredicate(info NodeInfo) algorithm.FitPredicate {
	selector := &NodeSelector{
		info: info,
	}
	return selector.PodSelectorMatches
}

func PodMatchesNodeLabels(pod *api.Pod, node *api.Node) bool {
	if len(pod.Spec.NodeSelector) == 0 {
		return true
	}
	selector := labels.SelectorFromSet(pod.Spec.NodeSelector)
	return selector.Matches(labels.Set(node.Labels))
}

type NodeSelector struct {
	info NodeInfo
}

func (n *NodeSelector) PodSelectorMatches(pod *api.Pod, existingPods []*api.Pod, node string) (bool, error) {
	minion, err := n.info.GetNodeInfo(node)
	if err != nil {
		return false, err
	}
	return PodMatchesNodeLabels(pod, minion), nil
}

func PodFitsHost(pod *api.Pod, existingPods []*api.Pod, node string) (bool, error) {
	if len(pod.Spec.NodeName) == 0 {
		return true, nil
	}
	return pod.Spec.NodeName == node, nil
}

type NodeLabelChecker struct {
	info     NodeInfo
	labels   []string
	presence bool
}

func NewNodeLabelPredicate(info NodeInfo, labels []string, presence bool) algorithm.FitPredicate {
	labelChecker := &NodeLabelChecker{
		info:     info,
		labels:   labels,
		presence: presence,
	}
	return labelChecker.CheckNodeLabelPresence
}

// CheckNodeLabelPresence checks whether all of the specified labels exists on a minion or not, regardless of their value
// If "presence" is false, then returns false if any of the requested labels matches any of the minion's labels,
// otherwise returns true.
// If "presence" is true, then returns false if any of the requested labels does not match any of the minion's labels,
// otherwise returns true.
//
// Consider the cases where the minions are placed in regions/zones/racks and these are identified by labels
// In some cases, it is required that only minions that are part of ANY of the defined regions/zones/racks be selected
//
// Alternately, eliminating minions that have a certain label, regardless of value, is also useful
// A minion may have a label with "retiring" as key and the date as the value
// and it may be desirable to avoid scheduling new pods on this minion
func (n *NodeLabelChecker) CheckNodeLabelPresence(pod *api.Pod, existingPods []*api.Pod, node string) (bool, error) {
	var exists bool
	minion, err := n.info.GetNodeInfo(node)
	if err != nil {
		return false, err
	}
	minionLabels := labels.Set(minion.Labels)
	for _, label := range n.labels {
		exists = minionLabels.Has(label)
		if (exists && !n.presence) || (!exists && n.presence) {
			return false, nil
		}
	}
	return true, nil
}

type ServiceAffinity struct {
	podLister     algorithm.PodLister
	serviceLister algorithm.ServiceLister
	nodeInfo      NodeInfo
	labels        []string
}

func NewServiceAffinityPredicate(podLister algorithm.PodLister, serviceLister algorithm.ServiceLister, nodeInfo NodeInfo, labels []string) algorithm.FitPredicate {
	affinity := &ServiceAffinity{
		podLister:     podLister,
		serviceLister: serviceLister,
		nodeInfo:      nodeInfo,
		labels:        labels,
	}
	return affinity.CheckServiceAffinity
}

// CheckServiceAffinity ensures that only the minions that match the specified labels are considered for scheduling.
// The set of labels to be considered are provided to the struct (ServiceAffinity).
// The pod is checked for the labels and any missing labels are then checked in the minion
// that hosts the service pods (peers) for the given pod.
//
// We add an implicit selector requiring some particular value V for label L to a pod, if:
// - L is listed in the ServiceAffinity object that is passed into the function
// - the pod does not have any NodeSelector for L
// - some other pod from the same service is already scheduled onto a minion that has value V for label L
func (s *ServiceAffinity) CheckServiceAffinity(pod *api.Pod, existingPods []*api.Pod, node string) (bool, error) {
	var affinitySelector labels.Selector

	// check if the pod being scheduled has the affinity labels specified in its NodeSelector
	affinityLabels := map[string]string{}
	nodeSelector := labels.Set(pod.Spec.NodeSelector)
	labelsExist := true
	for _, l := range s.labels {
		if nodeSelector.Has(l) {
			affinityLabels[l] = nodeSelector.Get(l)
		} else {
			// the current pod does not specify all the labels, look in the existing service pods
			labelsExist = false
		}
	}

	// skip looking at other pods in the service if the current pod defines all the required affinity labels
	if !labelsExist {
		services, err := s.serviceLister.GetPodServices(pod)
		if err == nil {
			// just use the first service and get the other pods within the service
			// TODO: a separate predicate can be created that tries to handle all services for the pod
			selector := labels.SelectorFromSet(services[0].Spec.Selector)
			servicePods, err := s.podLister.List(selector)
			if err != nil {
				return false, err
			}
			// consider only the pods that belong to the same namespace
			nsServicePods := []*api.Pod{}
			for _, nsPod := range servicePods {
				if nsPod.Namespace == pod.Namespace {
					nsServicePods = append(nsServicePods, nsPod)
				}
			}
			if len(nsServicePods) > 0 {
				// consider any service pod and fetch the minion its hosted on
				otherMinion, err := s.nodeInfo.GetNodeInfo(nsServicePods[0].Spec.NodeName)
				if err != nil {
					return false, err
				}
				for _, l := range s.labels {
					// If the pod being scheduled has the label value specified, do not override it
					if _, exists := affinityLabels[l]; exists {
						continue
					}
					if labels.Set(otherMinion.Labels).Has(l) {
						affinityLabels[l] = labels.Set(otherMinion.Labels).Get(l)
					}
				}
			}
		}
	}

	// if there are no existing pods in the service, consider all minions
	if len(affinityLabels) == 0 {
		affinitySelector = labels.Everything()
	} else {
		affinitySelector = labels.Set(affinityLabels).AsSelector()
	}

	minion, err := s.nodeInfo.GetNodeInfo(node)
	if err != nil {
		return false, err
	}

	// check if the minion matches the selector
	return affinitySelector.Matches(labels.Set(minion.Labels)), nil
}

func PodFitsPorts(pod *api.Pod, existingPods []*api.Pod, node string) (bool, error) {
	existingPorts := getUsedPorts(existingPods...)
	wantPorts := getUsedPorts(pod)
	for wport := range wantPorts {
		if wport == 0 {
			continue
		}
		if existingPorts[wport] {
			return false, nil
		}
	}
	return true, nil
}

func getUsedPorts(pods ...*api.Pod) map[int]bool {
	ports := make(map[int]bool)
	for _, pod := range pods {
		for _, container := range pod.Spec.Containers {
			for _, podPort := range container.Ports {
				ports[podPort.HostPort] = true
			}
		}
	}
	return ports
}

func filterNonRunningPods(pods []*api.Pod) []*api.Pod {
	if len(pods) == 0 {
		return pods
	}
	result := []*api.Pod{}
	for _, pod := range pods {
		if pod.Status.Phase == api.PodSucceeded || pod.Status.Phase == api.PodFailed {
			continue
		}
		result = append(result, pod)
	}
	return result
}

// MapPodsToMachines obtains a list of pods and pivots that list into a map where the keys are host names
// and the values are the list of pods running on that host.
func MapPodsToMachines(lister algorithm.PodLister) (map[string][]*api.Pod, error) {
	machineToPods := map[string][]*api.Pod{}
	// TODO: perform more targeted query...
	pods, err := lister.List(labels.Everything())
	if err != nil {
		return map[string][]*api.Pod{}, err
	}
	pods = filterNonRunningPods(pods)
	for _, scheduledPod := range pods {
		host := scheduledPod.Spec.NodeName
		machineToPods[host] = append(machineToPods[host], scheduledPod)
	}
	return machineToPods, nil
}
<commit_msg>Add more logging to scheduler predicates to help debugging max_pods e2e test flakyness<commit_after>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

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

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

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

package predicates

import (
	"fmt"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/plugin/pkg/scheduler/algorithm"

	"github.com/golang/glog"
)

type NodeInfo interface {
	GetNodeInfo(nodeID string) (*api.Node, error)
}

type StaticNodeInfo struct {
	*api.NodeList
}

func (nodes StaticNodeInfo) GetNodeInfo(nodeID string) (*api.Node, error) {
	for ix := range nodes.Items {
		if nodes.Items[ix].Name == nodeID {
			return &nodes.Items[ix], nil
		}
	}
	return nil, fmt.Errorf("failed to find node: %s, %#v", nodeID, nodes)
}

type ClientNodeInfo struct {
	*client.Client
}

func (nodes ClientNodeInfo) GetNodeInfo(nodeID string) (*api.Node, error) {
	return nodes.Nodes().Get(nodeID)
}

func isVolumeConflict(volume api.Volume, pod *api.Pod) bool {
	if volume.GCEPersistentDisk != nil {
		disk := volume.GCEPersistentDisk

		manifest := &(pod.Spec)
		for ix := range manifest.Volumes {
			if manifest.Volumes[ix].GCEPersistentDisk != nil &&
				manifest.Volumes[ix].GCEPersistentDisk.PDName == disk.PDName &&
				!(manifest.Volumes[ix].GCEPersistentDisk.ReadOnly && disk.ReadOnly) {
				return true
			}
		}
	}
	if volume.AWSElasticBlockStore != nil {
		volumeID := volume.AWSElasticBlockStore.VolumeID

		manifest := &(pod.Spec)
		for ix := range manifest.Volumes {
			if manifest.Volumes[ix].AWSElasticBlockStore != nil &&
				manifest.Volumes[ix].AWSElasticBlockStore.VolumeID == volumeID {
				return true
			}
		}
	}
	return false
}

// NoDiskConflict evaluates if a pod can fit due to the volumes it requests, and those that
// are already mounted. Some times of volumes are mounted onto node machines.  For now, these mounts
// are exclusive so if there is already a volume mounted on that node, another pod can't schedule
// there. This is GCE specific for now.
// TODO: migrate this into some per-volume specific code?
func NoDiskConflict(pod *api.Pod, existingPods []*api.Pod, node string) (bool, error) {
	manifest := &(pod.Spec)
	for ix := range manifest.Volumes {
		for podIx := range existingPods {
			if isVolumeConflict(manifest.Volumes[ix], existingPods[podIx]) {
				return false, nil
			}
		}
	}
	return true, nil
}

type ResourceFit struct {
	info NodeInfo
}

type resourceRequest struct {
	milliCPU int64
	memory   int64
}

func getResourceRequest(pod *api.Pod) resourceRequest {
	result := resourceRequest{}
	for ix := range pod.Spec.Containers {
		limits := pod.Spec.Containers[ix].Resources.Limits
		result.memory += limits.Memory().Value()
		result.milliCPU += limits.Cpu().MilliValue()
	}
	return result
}

func CheckPodsExceedingCapacity(pods []*api.Pod, capacity api.ResourceList) (fitting []*api.Pod, notFitting []*api.Pod) {
	totalMilliCPU := capacity.Cpu().MilliValue()
	totalMemory := capacity.Memory().Value()
	milliCPURequested := int64(0)
	memoryRequested := int64(0)
	for _, pod := range pods {
		podRequest := getResourceRequest(pod)
		fitsCPU := totalMilliCPU == 0 || (totalMilliCPU-milliCPURequested) >= podRequest.milliCPU
		fitsMemory := totalMemory == 0 || (totalMemory-memoryRequested) >= podRequest.memory
		if !fitsCPU || !fitsMemory {
			// the pod doesn't fit
			notFitting = append(notFitting, pod)
			continue
		}
		// the pod fits
		milliCPURequested += podRequest.milliCPU
		memoryRequested += podRequest.memory
		fitting = append(fitting, pod)
	}
	return
}

// PodFitsResources calculates fit based on requested, rather than used resources
func (r *ResourceFit) PodFitsResources(pod *api.Pod, existingPods []*api.Pod, node string) (bool, error) {
	podRequest := getResourceRequest(pod)
	info, err := r.info.GetNodeInfo(node)
	if err != nil {
		return false, err
	}
	if podRequest.milliCPU == 0 && podRequest.memory == 0 {
		return int64(len(existingPods)) < info.Status.Capacity.Pods().Value(), nil
	}
	pods := []*api.Pod{}
	copy(pods, existingPods)
	pods = append(existingPods, pod)
	_, exceeding := CheckPodsExceedingCapacity(pods, info.Status.Capacity)
	if len(exceeding) > 0 || int64(len(pods)) > info.Status.Capacity.Pods().Value() {
		glog.V(4).Infof("Cannot schedule Pod %v, because Node %v is full, running %v out of %v Pods.", pod, node, len(pods)-1, info.Status.Capacity.Pods().Value())
		return false, nil
	}
	glog.V(4).Infof("Schedule Pod %v on Node %v is allowed, Node is running only %v out of %v Pods.", pod, node, len(pods)-1, info.Status.Capacity.Pods().Value())
	return true, nil
}

func NewResourceFitPredicate(info NodeInfo) algorithm.FitPredicate {
	fit := &ResourceFit{
		info: info,
	}
	return fit.PodFitsResources
}

func NewSelectorMatchPredicate(info NodeInfo) algorithm.FitPredicate {
	selector := &NodeSelector{
		info: info,
	}
	return selector.PodSelectorMatches
}

func PodMatchesNodeLabels(pod *api.Pod, node *api.Node) bool {
	if len(pod.Spec.NodeSelector) == 0 {
		return true
	}
	selector := labels.SelectorFromSet(pod.Spec.NodeSelector)
	return selector.Matches(labels.Set(node.Labels))
}

type NodeSelector struct {
	info NodeInfo
}

func (n *NodeSelector) PodSelectorMatches(pod *api.Pod, existingPods []*api.Pod, node string) (bool, error) {
	minion, err := n.info.GetNodeInfo(node)
	if err != nil {
		return false, err
	}
	return PodMatchesNodeLabels(pod, minion), nil
}

func PodFitsHost(pod *api.Pod, existingPods []*api.Pod, node string) (bool, error) {
	if len(pod.Spec.NodeName) == 0 {
		return true, nil
	}
	return pod.Spec.NodeName == node, nil
}

type NodeLabelChecker struct {
	info     NodeInfo
	labels   []string
	presence bool
}

func NewNodeLabelPredicate(info NodeInfo, labels []string, presence bool) algorithm.FitPredicate {
	labelChecker := &NodeLabelChecker{
		info:     info,
		labels:   labels,
		presence: presence,
	}
	return labelChecker.CheckNodeLabelPresence
}

// CheckNodeLabelPresence checks whether all of the specified labels exists on a minion or not, regardless of their value
// If "presence" is false, then returns false if any of the requested labels matches any of the minion's labels,
// otherwise returns true.
// If "presence" is true, then returns false if any of the requested labels does not match any of the minion's labels,
// otherwise returns true.
//
// Consider the cases where the minions are placed in regions/zones/racks and these are identified by labels
// In some cases, it is required that only minions that are part of ANY of the defined regions/zones/racks be selected
//
// Alternately, eliminating minions that have a certain label, regardless of value, is also useful
// A minion may have a label with "retiring" as key and the date as the value
// and it may be desirable to avoid scheduling new pods on this minion
func (n *NodeLabelChecker) CheckNodeLabelPresence(pod *api.Pod, existingPods []*api.Pod, node string) (bool, error) {
	var exists bool
	minion, err := n.info.GetNodeInfo(node)
	if err != nil {
		return false, err
	}
	minionLabels := labels.Set(minion.Labels)
	for _, label := range n.labels {
		exists = minionLabels.Has(label)
		if (exists && !n.presence) || (!exists && n.presence) {
			return false, nil
		}
	}
	return true, nil
}

type ServiceAffinity struct {
	podLister     algorithm.PodLister
	serviceLister algorithm.ServiceLister
	nodeInfo      NodeInfo
	labels        []string
}

func NewServiceAffinityPredicate(podLister algorithm.PodLister, serviceLister algorithm.ServiceLister, nodeInfo NodeInfo, labels []string) algorithm.FitPredicate {
	affinity := &ServiceAffinity{
		podLister:     podLister,
		serviceLister: serviceLister,
		nodeInfo:      nodeInfo,
		labels:        labels,
	}
	return affinity.CheckServiceAffinity
}

// CheckServiceAffinity ensures that only the minions that match the specified labels are considered for scheduling.
// The set of labels to be considered are provided to the struct (ServiceAffinity).
// The pod is checked for the labels and any missing labels are then checked in the minion
// that hosts the service pods (peers) for the given pod.
//
// We add an implicit selector requiring some particular value V for label L to a pod, if:
// - L is listed in the ServiceAffinity object that is passed into the function
// - the pod does not have any NodeSelector for L
// - some other pod from the same service is already scheduled onto a minion that has value V for label L
func (s *ServiceAffinity) CheckServiceAffinity(pod *api.Pod, existingPods []*api.Pod, node string) (bool, error) {
	var affinitySelector labels.Selector

	// check if the pod being scheduled has the affinity labels specified in its NodeSelector
	affinityLabels := map[string]string{}
	nodeSelector := labels.Set(pod.Spec.NodeSelector)
	labelsExist := true
	for _, l := range s.labels {
		if nodeSelector.Has(l) {
			affinityLabels[l] = nodeSelector.Get(l)
		} else {
			// the current pod does not specify all the labels, look in the existing service pods
			labelsExist = false
		}
	}

	// skip looking at other pods in the service if the current pod defines all the required affinity labels
	if !labelsExist {
		services, err := s.serviceLister.GetPodServices(pod)
		if err == nil {
			// just use the first service and get the other pods within the service
			// TODO: a separate predicate can be created that tries to handle all services for the pod
			selector := labels.SelectorFromSet(services[0].Spec.Selector)
			servicePods, err := s.podLister.List(selector)
			if err != nil {
				return false, err
			}
			// consider only the pods that belong to the same namespace
			nsServicePods := []*api.Pod{}
			for _, nsPod := range servicePods {
				if nsPod.Namespace == pod.Namespace {
					nsServicePods = append(nsServicePods, nsPod)
				}
			}
			if len(nsServicePods) > 0 {
				// consider any service pod and fetch the minion its hosted on
				otherMinion, err := s.nodeInfo.GetNodeInfo(nsServicePods[0].Spec.NodeName)
				if err != nil {
					return false, err
				}
				for _, l := range s.labels {
					// If the pod being scheduled has the label value specified, do not override it
					if _, exists := affinityLabels[l]; exists {
						continue
					}
					if labels.Set(otherMinion.Labels).Has(l) {
						affinityLabels[l] = labels.Set(otherMinion.Labels).Get(l)
					}
				}
			}
		}
	}

	// if there are no existing pods in the service, consider all minions
	if len(affinityLabels) == 0 {
		affinitySelector = labels.Everything()
	} else {
		affinitySelector = labels.Set(affinityLabels).AsSelector()
	}

	minion, err := s.nodeInfo.GetNodeInfo(node)
	if err != nil {
		return false, err
	}

	// check if the minion matches the selector
	return affinitySelector.Matches(labels.Set(minion.Labels)), nil
}

func PodFitsPorts(pod *api.Pod, existingPods []*api.Pod, node string) (bool, error) {
	existingPorts := getUsedPorts(existingPods...)
	wantPorts := getUsedPorts(pod)
	for wport := range wantPorts {
		if wport == 0 {
			continue
		}
		if existingPorts[wport] {
			return false, nil
		}
	}
	return true, nil
}

func getUsedPorts(pods ...*api.Pod) map[int]bool {
	ports := make(map[int]bool)
	for _, pod := range pods {
		for _, container := range pod.Spec.Containers {
			for _, podPort := range container.Ports {
				ports[podPort.HostPort] = true
			}
		}
	}
	return ports
}

func filterNonRunningPods(pods []*api.Pod) []*api.Pod {
	if len(pods) == 0 {
		return pods
	}
	result := []*api.Pod{}
	for _, pod := range pods {
		if pod.Status.Phase == api.PodSucceeded || pod.Status.Phase == api.PodFailed {
			continue
		}
		result = append(result, pod)
	}
	return result
}

// MapPodsToMachines obtains a list of pods and pivots that list into a map where the keys are host names
// and the values are the list of pods running on that host.
func MapPodsToMachines(lister algorithm.PodLister) (map[string][]*api.Pod, error) {
	machineToPods := map[string][]*api.Pod{}
	// TODO: perform more targeted query...
	pods, err := lister.List(labels.Everything())
	if err != nil {
		return map[string][]*api.Pod{}, err
	}
	pods = filterNonRunningPods(pods)
	for _, scheduledPod := range pods {
		host := scheduledPod.Spec.NodeName
		machineToPods[host] = append(machineToPods[host], scheduledPod)
	}
	return machineToPods, nil
}
<|endoftext|>


<commit_before>package filesystem_test

import (
	"errors"
	"fmt"
	"os"
	"reflect"
	"testing"

	"github.com/naoty/todo/repository/filesystem"
	"github.com/naoty/todo/todo"
)

func TestGet(t *testing.T) {
	testcases := []struct {
		input  int
		output *todo.Todo
		err    error
	}{
		{1, &todo.Todo{ID: 1, Title: "dummy", State: todo.Undone}, nil},
		{2, &todo.Todo{ID: 2, Title: "dummy", State: todo.Done}, nil},
		{3, &todo.Todo{ID: 3, Title: "dummy", State: todo.Waiting}, nil},
		{1000, nil, filesystem.ErrTODONotFound},
	}

	repo, err := filesystem.New("./testdata/todos")
	if err != nil {
		t.Fatalf("failed to initialize repository: %v", err)
	}

	for _, testcase := range testcases {
		name := fmt.Sprintf("ID:%d", testcase.input)
		t.Run(name, func(t *testing.T) {
			td, err := repo.Get(testcase.input)

			if err != nil {
				if !errors.Is(err, testcase.err) {
					t.Errorf("got: %v, want: %v", err, testcase.err)
				}
				return
			}

			if !td.Equal(testcase.output) {
				t.Errorf("got: %+v, want: %+v", td, testcase.output)
			}
		})
	}
}

func TestList(t *testing.T) {
	repo, err := filesystem.New("./testdata/todos")
	if err != nil {
		t.Fatalf("failed to initialize repository: %v", err)
	}

	todos, err := repo.List()
	if err != nil {
		t.Fatalf("failed to list todos")
	}

	ids := make([]int, 3)
	for i, td := range todos {
		ids[i] = td.ID
	}

	want := []int{2, 1, 3}
	if !reflect.DeepEqual(ids, want) {
		t.Errorf("got: %v, want: %v", ids, want)
	}
}

func TestAdd(t *testing.T) {
	repo, err := filesystem.New("./testdata/sandbox")
	if err != nil {
		t.Fatalf("failed to initialize repository: %v", err)
	}

	t.Cleanup(func() {
		err := os.RemoveAll("./testdata/sandbox")
		if err != nil {
			t.Fatalf("failed to cleanup sandbox: %v", err)
		}
	})

	parent := 0
	err = repo.Add("dummy", &parent)
	if err != nil {
		t.Fatalf("failed to add a TODO")
	}

	td, err := repo.Get(1)
	if err != nil {
		t.Fatalf("failed to get a new TODO")
	}

	if td.Title != "dummy" {
		t.Errorf("got: %s, want: dummy", td.Title)
	}

	if td.State != todo.Undone {
		t.Errorf("got: %s, want: %s", td.State, todo.Undone)
	}

	if td.Body != "" {
		t.Errorf("got: %s, want: ''", td.Body)
	}
}

func TestUpdate(t *testing.T) {
	repo, err := filesystem.New("./testdata/sandbox")
	if err != nil {
		t.Fatalf("failed to initialize repository: %v", err)
	}

	t.Cleanup(func() {
		err := os.RemoveAll("./testdata/sandbox")
		if err != nil {
			t.Fatalf("failed to cleanup sandbox: %v", err)
		}
	})

	parent := 0
	err = repo.Add("dummy", &parent)
	if err != nil {
		t.Fatal("failed to add a TODO")
	}

	td, err := repo.Get(1)
	if err != nil {
		t.Fatal("failed to get a new TODO")
	}

	td.State = todo.Done
	err = repo.Update(td)
	if err != nil {
		t.Fatal("failed to update a TODO")
	}

	td, err = repo.Get(1)
	if err != nil {
		t.Fatal("failed to get a new TODO")
	}

	if td.State != todo.Done {
		t.Errorf("got: %s, want: %s", td.State, todo.Done)
	}
}
<commit_msg>Add a test for (*FileSystem).Delete()<commit_after>package filesystem_test

import (
	"errors"
	"fmt"
	"os"
	"reflect"
	"testing"

	"github.com/naoty/todo/repository/filesystem"
	"github.com/naoty/todo/todo"
)

func TestGet(t *testing.T) {
	testcases := []struct {
		input  int
		output *todo.Todo
		err    error
	}{
		{1, &todo.Todo{ID: 1, Title: "dummy", State: todo.Undone}, nil},
		{2, &todo.Todo{ID: 2, Title: "dummy", State: todo.Done}, nil},
		{3, &todo.Todo{ID: 3, Title: "dummy", State: todo.Waiting}, nil},
		{1000, nil, filesystem.ErrTODONotFound},
	}

	repo, err := filesystem.New("./testdata/todos")
	if err != nil {
		t.Fatalf("failed to initialize repository: %v", err)
	}

	for _, testcase := range testcases {
		name := fmt.Sprintf("ID:%d", testcase.input)
		t.Run(name, func(t *testing.T) {
			td, err := repo.Get(testcase.input)

			if err != nil {
				if !errors.Is(err, testcase.err) {
					t.Errorf("got: %v, want: %v", err, testcase.err)
				}
				return
			}

			if !td.Equal(testcase.output) {
				t.Errorf("got: %+v, want: %+v", td, testcase.output)
			}
		})
	}
}

func TestList(t *testing.T) {
	repo, err := filesystem.New("./testdata/todos")
	if err != nil {
		t.Fatalf("failed to initialize repository: %v", err)
	}

	todos, err := repo.List()
	if err != nil {
		t.Fatalf("failed to list todos")
	}

	ids := make([]int, 3)
	for i, td := range todos {
		ids[i] = td.ID
	}

	want := []int{2, 1, 3}
	if !reflect.DeepEqual(ids, want) {
		t.Errorf("got: %v, want: %v", ids, want)
	}
}

func TestAdd(t *testing.T) {
	repo, err := filesystem.New("./testdata/sandbox")
	if err != nil {
		t.Fatalf("failed to initialize repository: %v", err)
	}

	t.Cleanup(func() {
		err := os.RemoveAll("./testdata/sandbox")
		if err != nil {
			t.Fatalf("failed to cleanup sandbox: %v", err)
		}
	})

	parent := 0
	err = repo.Add("dummy", &parent)
	if err != nil {
		t.Fatalf("failed to add a TODO")
	}

	td, err := repo.Get(1)
	if err != nil {
		t.Fatalf("failed to get a new TODO")
	}

	if td.Title != "dummy" {
		t.Errorf("got: %s, want: dummy", td.Title)
	}

	if td.State != todo.Undone {
		t.Errorf("got: %s, want: %s", td.State, todo.Undone)
	}

	if td.Body != "" {
		t.Errorf("got: %s, want: ''", td.Body)
	}
}

func TestUpdate(t *testing.T) {
	repo, err := filesystem.New("./testdata/sandbox")
	if err != nil {
		t.Fatalf("failed to initialize repository: %v", err)
	}

	t.Cleanup(func() {
		err := os.RemoveAll("./testdata/sandbox")
		if err != nil {
			t.Fatalf("failed to cleanup sandbox: %v", err)
		}
	})

	parent := 0
	err = repo.Add("dummy", &parent)
	if err != nil {
		t.Fatal("failed to add a TODO")
	}

	td, err := repo.Get(1)
	if err != nil {
		t.Fatal("failed to get a new TODO")
	}

	td.State = todo.Done
	err = repo.Update(td)
	if err != nil {
		t.Fatal("failed to update a TODO")
	}

	td, err = repo.Get(1)
	if err != nil {
		t.Fatal("failed to get a new TODO")
	}

	if td.State != todo.Done {
		t.Errorf("got: %s, want: %s", td.State, todo.Done)
	}
}

func TestDelete(t *testing.T) {
	repo, err := filesystem.New("./testdata/sandbox")
	if err != nil {
		t.Fatalf("failed to initialize repository: %v", err)
	}

	t.Cleanup(func() {
		err := os.RemoveAll("./testdata/sandbox")
		if err != nil {
			t.Fatalf("failed to cleanup sandbox: %v", err)
		}
	})

	parent := 0
	err = repo.Add("dummy", &parent)
	if err != nil {
		t.Fatal("failed to add a TODO")
	}

	err = repo.Delete(1)
	if err != nil {
		t.Fatal("failed to delete a TODO")
	}

	if _, err := os.Stat("./testdata/sandbox/1.md"); os.IsExist(err) {
		t.Error("1.md is not deleted")
	}
}
<|endoftext|>


<commit_before>package analyze

import (
	"bytes"
	"context"
	"fmt"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/bazelbuild/buildtools/edit"
	"github.com/bazelbuild/rules_typescript/ts_auto_deps/workspace"
	"github.com/golang/protobuf/proto"

	appb "github.com/bazelbuild/buildtools/build_proto"
)

// pkgCacheEntry represents a set of loaded rules and a mapping from alias
// to rules from a package.
type pkgCacheEntry struct {
	// rules is all rules in a package.
	rules []*appb.Rule
	// aliases is a map from an alias label to the actual rule of the alias.
	aliases map[string]*appb.Rule
}

// QueryBasedTargetLoader uses Bazel query to load targets from BUILD files.
type QueryBasedTargetLoader struct {
	workdir     string
	bazelBinary string

	// pkgCache is a mapping from a package to all of the rules in said
	// package along with a map from aliases to actual rules.
	//
	// Keys are of the form of "<visibility>|<package>" where visibility
	// is the package that rules in package must be visible to and package
	// is the actual package that has been loaded and cached.
	//
	// Since a new target loader is constructed for each directory being
	// analyzed in the "-recursive" case, these caches will be garbage
	// collected between directories.
	pkgCache map[string]*pkgCacheEntry
	// labelCache is a mapping from a label to its loaded rule.
	labelCache map[string]*appb.Rule

	// queryCount is the total number of queries executed by the target loader.
	queryCount int
}

// NewQueryBasedTargetLoader constructs a new QueryBasedTargetLoader rooted
// in workdir.
func NewQueryBasedTargetLoader(workdir, bazelBinary string) *QueryBasedTargetLoader {
	return &QueryBasedTargetLoader{
		workdir:     workdir,
		bazelBinary: bazelBinary,

		pkgCache:   make(map[string]*pkgCacheEntry),
		labelCache: make(map[string]*appb.Rule),
	}
}

// LoadLabels uses Bazel query to load targets associated with labels from BUILD
// files.
func (q *QueryBasedTargetLoader) LoadLabels(pkg string, labels []string) (map[string]*appb.Rule, error) {
	var labelCacheMisses []string
	for _, label := range labels {
		if _, ok := q.labelCache[labelCacheKey(pkg, label)]; !ok {
			labelCacheMisses = append(labelCacheMisses, label)
		}
	}
	if len(labelCacheMisses) > 0 {
		var queries []string
		if pkg == "" {
			queries = labelCacheMisses
		} else {
			for _, label := range labelCacheMisses {
				queries = append(queries, fmt.Sprintf("visible(%s:*, %s)", pkg, label))
			}
		}
		r, err := q.batchQuery(queries)
		if err != nil {
			return nil, err
		}
		for _, target := range r.GetTarget() {
			label, err := q.ruleLabel(target)
			if err != nil {
				return nil, err
			}
			q.labelCache[labelCacheKey(pkg, label)] = target.GetRule()
		}
		for _, label := range labelCacheMisses {
			key := labelCacheKey(pkg, label)
			if _, ok := q.labelCache[key]; !ok {
				// Set to nil so the result exists in the cache and is not
				// loaded again. If the nil is not added at the appropriate
				// cache key, LoadLabels will attempt to load it again when
				// next requested instead of getting a cache hit.
				q.labelCache[key] = nil
			}
		}
	}
	labelToRule := make(map[string]*appb.Rule)
	for _, label := range labels {
		labelToRule[label] = q.labelCache[labelCacheKey(pkg, label)]
	}
	return labelToRule, nil
}

func labelCacheKey(currentPkg, label string) string {
	return currentPkg + "^" + label
}

// LoadImportPaths uses Bazel Query to load targets associated with import
// paths from BUILD files.
func (q *QueryBasedTargetLoader) LoadImportPaths(ctx context.Context, currentPkg, workspaceRoot string, paths []string) (map[string]*appb.Rule, error) {
	debugf("loading imports visible to %q relative to %q: %q", currentPkg, workspaceRoot, paths)
	results := make(map[string]*appb.Rule)

	addedPaths := make(map[string]bool)
	var possiblePaths []string
	for _, path := range paths {
		if strings.HasPrefix(path, "goog:") {
			// 'goog:' imports are resolved using an sstable.
			results[path] = nil
			continue
		}
		if !strings.HasPrefix(path, "@") {
			if _, ok := addedPaths[path]; !ok {
				addedPaths[path] = true

				// If the path has a suffix of ".ngfactory" or ".ngsummary", it might
				// be an Angular AOT generated file. We can infer the target as we
				// infer its corresponding ngmodule target by simply stripping the
				// ".ngfactory" / ".ngsummary" suffix
				path = strings.TrimSuffix(strings.TrimSuffix(path, ".ngsummary"), ".ngfactory")
				path = strings.TrimPrefix(path, workspace.Name()+"/")

				possiblePaths = append(possiblePaths, pathWithExtensions(path)...)
				possiblePaths = append(possiblePaths, pathWithExtensions(filepath.Join(path, "index"))...)
			}
		}
	}

	r, err := q.batchQuery(possiblePaths)
	if err != nil {
		return nil, err
	}
	var sourceFileLabels, generators []string
	generatorsToFiles := make(map[string][]*appb.GeneratedFile)
	for _, target := range r.GetTarget() {
		label, err := q.fileLabel(target)
		if err != nil {
			return nil, err
		}
		switch target.GetType() {
		case appb.Target_GENERATED_FILE:
			file := target.GetGeneratedFile()
			generator := file.GetGeneratingRule()

			generators = append(generators, generator)
			generatorsToFiles[generator] = append(generatorsToFiles[generator], file)
		case appb.Target_SOURCE_FILE:
			sourceFileLabels = append(sourceFileLabels, label)
		}
	}

	labelToRule := make(map[string]*appb.Rule)
	for len(generators) > 0 {
		generatorToRule, err := q.LoadLabels(currentPkg, generators)
		if err != nil {
			return nil, err
		}
		var newGenerators []string
		for label, rule := range generatorToRule {
			_, _, target := edit.ParseLabel(label)
			if generator := stringAttribute(rule, "generator_name"); generator != "" && generator != target {
				// Located rule is also a generated rule. Look for the rule
				// that generates it.
				_, pkg, _ := edit.ParseLabel(label)
				newLabel := "//" + pkg + ":" + generator
				newGenerators = append(newGenerators, newLabel)
				generatorsToFiles[newLabel] = generatorsToFiles[label]
			} else {
				for _, generated := range generatorsToFiles[label] {
					labelToRule[generated.GetName()] = rule
				}
			}
		}
		generators = newGenerators
	}

	sourceLabelToRule, err := q.loadRulesIncludingSourceFiles(workspaceRoot, sourceFileLabels)
	if err != nil {
		return nil, err
	}
	for label, rule := range sourceLabelToRule {
		labelToRule[label] = rule
	}

	for label, rule := range labelToRule {
		_, pkg, file := edit.ParseLabel(label)
		// Trim "/index" suffixes that were added to path in the queries above.
		pathWithoutExtension := strings.TrimSuffix(filepath.Join(pkg, stripTSExtension(file)), string(filepath.Separator)+"index")
		for _, path := range paths {
			if pathWithoutExtension == strings.TrimSuffix(path, string(filepath.Separator)+"index") {
				results[path] = rule
			} else if pathWithoutExtension == strings.TrimSuffix(path, ".ngsummary") {
				results[path] = rule
			} else if pathWithoutExtension == strings.TrimSuffix(path, ".ngfactory") {
				results[path] = rule
			}
		}
	}

	return results, nil
}

func (q *QueryBasedTargetLoader) ruleLabel(target *appb.Target) (string, error) {
	if t := target.GetType(); t != appb.Target_RULE {
		return "", fmt.Errorf("target contains object of type %q instead of type %q", t, appb.Target_RULE)
	}
	return target.GetRule().GetName(), nil
}

func (q *QueryBasedTargetLoader) fileLabel(target *appb.Target) (string, error) {
	switch t := target.GetType(); t {
	case appb.Target_GENERATED_FILE:
		return target.GetGeneratedFile().GetName(), nil
	case appb.Target_SOURCE_FILE:
		return target.GetSourceFile().GetName(), nil
	default:
		return "", fmt.Errorf("target contains object of type %q instead of type %q or %q", t, appb.Target_SOURCE_FILE, appb.Target_GENERATED_FILE)
	}
}

// loadRuleIncludingSourceFiles loads all rules which include labels in
// sourceFileLabels, Returns a map from source file label to the rule which
// includes it.
func (q *QueryBasedTargetLoader) loadRulesIncludingSourceFiles(workspaceRoot string, sourceFileLabels []string) (map[string]*appb.Rule, error) {
	pkgToLabels := make(map[string][]string)
	queries := make([]string, 0, len(sourceFileLabels))
	for _, label := range sourceFileLabels {
		_, pkg, file := edit.ParseLabel(label)
		pkgToLabels[pkg] = append(pkgToLabels[pkg], label)
		// Query for all targets in the package which use file.
		queries = append(queries, fmt.Sprintf("attr('srcs', %s, //%s:*)", file, pkg))
	}
	r, err := q.batchQuery(queries)
	if err != nil {
		return nil, err
	}
	labelToRule := make(map[string]*appb.Rule)
	for _, target := range r.GetTarget() {
		label, err := q.ruleLabel(target)
		if err != nil {
			return nil, err
		}
		rule := target.GetRule()
		_, pkg, _ := edit.ParseLabel(label)
		labels := pkgToLabels[pkg]
		for _, src := range listAttribute(rule, "srcs") {
			for _, l := range labels {
				if src == l {
					labelToRule[l] = rule
					break
				}
			}
		}
	}
	return labelToRule, nil
}

// batchQuery runs a set of queries with a single call to Bazel query and the
// '--keep_going' flag.
func (q *QueryBasedTargetLoader) batchQuery(queries []string) (*appb.QueryResult, error) {
	// Join all of the queries with a '+' character according to Bazel's
	// syntax for running multiple queries.
	return q.query("--keep_going", strings.Join(queries, "+"))
}

func (q *QueryBasedTargetLoader) query(args ...string) (*appb.QueryResult, error) {
	n := len(args)
	if n < 1 {
		return nil, fmt.Errorf("expected at least one argument")
	}
	query := args[n-1]
	if query == "" {
		// An empty query was provided so return an empty result without
		// making a call to Bazel.
		return &appb.QueryResult{}, nil
	}
	var stdout, stderr bytes.Buffer
	args = append([]string{"query", "--output=proto"}, args...)
	q.queryCount++
	debugf("executing query #%d in %q: %s %s %q", q.queryCount, q.workdir, q.bazelBinary, strings.Join(args[:len(args)-1], " "), query)
	cmd := exec.Command(q.bazelBinary, args...)
	cmd.Dir = q.workdir
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	startTime := time.Now()
	if err := cmd.Run(); err != nil {
		// Exit status 3 is a direct result of one or more queries in a set of
		// queries not returning a result while running with the '--keep_going'
		// flag. Since one query failing to return a result does not hinder the
		// other queries from returning a result, ignore these errors.
		if err.Error() != "exit status 3" {
			// The error provided as a result is less useful than the contents of
			// stderr for debugging.
			return nil, fmt.Errorf(stderr.String())
		}
	}
	debugf("query #%d took %v", q.queryCount, time.Since(startTime))
	var result appb.QueryResult
	if err := proto.Unmarshal(stdout.Bytes(), &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// dedupeLabels returns a new set of labels with no duplicates.
func dedupeLabels(labels []string) []string {
	addedLabels := make(map[string]bool)
	var uniqueLabels []string
	for _, label := range labels {
		if _, added := addedLabels[label]; !added {
			addedLabels[label] = true
			uniqueLabels = append(uniqueLabels, label)
		}
	}
	return uniqueLabels
}

// isTazeManagedRuleClass checks if a class is a ts_auto_deps-managed rule class.
func isTazeManagedRuleClass(class string) bool {
	for _, c := range []string{
		"ts_library",
		// TODO(alexeagle): Add ts_declaration once it can be determined
		// if they are unused.
		"ng_module",
		"js_library",
	} {
		if c == class {
			return true
		}
	}
	return false
}

// typeScriptRules returns all TypeScript rules in rules.
func typeScriptRules(rules []*appb.Rule) []*appb.Rule {
	var tsRules []*appb.Rule
	for _, rule := range rules {
		for _, supportedRuleClass := range []string{
			"ts_library",
			"ts_declaration",
			"ng_module",
		} {
			if rule.GetRuleClass() == supportedRuleClass {
				tsRules = append(tsRules, rule)
				break
			}
		}
	}
	return tsRules
}

// resolveAgainstModuleRoot resolves imported against moduleRoot and moduleName.
func resolveAgainstModuleRoot(label, moduleRoot, moduleName, imported string) string {
	if moduleRoot == "" && moduleName == "" {
		return imported
	}
	trim := strings.TrimPrefix(imported, moduleName)
	if trim == imported {
		return imported
	}
	_, pkg, _ := edit.ParseLabel(label)
	return filepath.Join(pkg, moduleRoot, trim)
}

// parsePackageName parses and returns the scope and package of imported. For
// example, "@foo/bar" would have a scope of "@foo" and a package of "bar".
func parsePackageName(imported string) (string, string) {
	firstSlash := strings.Index(imported, "/")
	if firstSlash == -1 {
		return imported, ""
	}
	afterSlash := imported[firstSlash+1:]
	if secondSlash := strings.Index(afterSlash, "/"); secondSlash > -1 {
		return imported[:firstSlash], afterSlash[:secondSlash]
	}
	return imported[:firstSlash], afterSlash
}
<commit_msg>No external change<commit_after>package analyze

import (
	"bytes"
	"context"
	"fmt"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/bazelbuild/buildtools/edit"
	"github.com/bazelbuild/rules_typescript/ts_auto_deps/workspace"
	"github.com/golang/protobuf/proto"

	appb "github.com/bazelbuild/buildtools/build_proto"
)

// pkgCacheEntry represents a set of loaded rules and a mapping from alias
// to rules from a package.
type pkgCacheEntry struct {
	// rules is all rules in a package.
	rules []*appb.Rule
	// aliases is a map from an alias label to the actual rule of the alias.
	aliases map[string]*appb.Rule
}

// QueryBasedTargetLoader uses Bazel query to load targets from BUILD files.
type QueryBasedTargetLoader struct {
	workdir     string
	bazelBinary string

	// pkgCache is a mapping from a package to all of the rules in said
	// package along with a map from aliases to actual rules.
	//
	// Keys are of the form of "<visibility>|<package>" where visibility
	// is the package that rules in package must be visible to and package
	// is the actual package that has been loaded and cached.
	//
	// Since a new target loader is constructed for each directory being
	// analyzed in the "-recursive" case, these caches will be garbage
	// collected between directories.
	pkgCache map[string]*pkgCacheEntry
	// labelCache is a mapping from a label to its loaded rule.
	labelCache map[string]*appb.Rule

	// queryCount is the total number of queries executed by the target loader.
	queryCount int
}

// NewQueryBasedTargetLoader constructs a new QueryBasedTargetLoader rooted
// in workdir.
func NewQueryBasedTargetLoader(workdir, bazelBinary string) *QueryBasedTargetLoader {
	return &QueryBasedTargetLoader{
		workdir:     workdir,
		bazelBinary: bazelBinary,

		pkgCache:   make(map[string]*pkgCacheEntry),
		labelCache: make(map[string]*appb.Rule),
	}
}

// LoadLabels uses Bazel query to load targets associated with labels from BUILD
// files.
func (q *QueryBasedTargetLoader) LoadLabels(pkg string, labels []string) (map[string]*appb.Rule, error) {
	var labelCacheMisses []string
	for _, label := range labels {
		if _, ok := q.labelCache[labelCacheKey(pkg, label)]; !ok {
			labelCacheMisses = append(labelCacheMisses, label)
		}
	}
	if len(labelCacheMisses) > 0 {
		var queries []string
		if pkg == "" {
			queries = labelCacheMisses
		} else {
			for _, label := range labelCacheMisses {
				queries = append(queries, fmt.Sprintf("visible(%s:*, %s)", pkg, label))
			}
		}
		r, err := q.batchQuery(queries)
		if err != nil {
			return nil, err
		}
		for _, target := range r.GetTarget() {
			label, err := q.ruleLabel(target)
			if err != nil {
				return nil, err
			}
			q.labelCache[labelCacheKey(pkg, label)] = target.GetRule()
		}
		for _, label := range labelCacheMisses {
			key := labelCacheKey(pkg, label)
			if _, ok := q.labelCache[key]; !ok {
				// Set to nil so the result exists in the cache and is not
				// loaded again. If the nil is not added at the appropriate
				// cache key, LoadLabels will attempt to load it again when
				// next requested instead of getting a cache hit.
				q.labelCache[key] = nil
			}
		}
	}
	labelToRule := make(map[string]*appb.Rule)
	for _, label := range labels {
		labelToRule[label] = q.labelCache[labelCacheKey(pkg, label)]
	}
	return labelToRule, nil
}

func labelCacheKey(currentPkg, label string) string {
	return currentPkg + "^" + label
}

// LoadImportPaths uses Bazel Query to load targets associated with import
// paths from BUILD files.
func (q *QueryBasedTargetLoader) LoadImportPaths(ctx context.Context, currentPkg, workspaceRoot string, paths []string) (map[string]*appb.Rule, error) {
	debugf("loading imports visible to %q relative to %q: %q", currentPkg, workspaceRoot, paths)
	results := make(map[string]*appb.Rule)

	addedPaths := make(map[string]bool)
	var possiblePaths []string
	for _, path := range paths {
		if strings.HasPrefix(path, "goog:") {
			// 'goog:' imports are resolved using an sstable.
			results[path] = nil
			continue
		}
		if !strings.HasPrefix(path, "@") {
			if _, ok := addedPaths[path]; !ok {
				addedPaths[path] = true

				// If the path has a suffix of ".ngfactory" or ".ngsummary", it might
				// be an Angular AOT generated file. We can infer the target as we
				// infer its corresponding ngmodule target by simply stripping the
				// ".ngfactory" / ".ngsummary" suffix
				path = strings.TrimSuffix(strings.TrimSuffix(path, ".ngsummary"), ".ngfactory")
				path = strings.TrimPrefix(path, workspace.Name()+"/")

				possiblePaths = append(possiblePaths, pathWithExtensions(path)...)
				possiblePaths = append(possiblePaths, pathWithExtensions(filepath.Join(path, "index"))...)
			}
		}
	}

	r, err := q.batchQuery(possiblePaths)
	if err != nil {
		return nil, err
	}
	var sourceFileLabels, generators []string
	generatorsToFiles := make(map[string][]*appb.GeneratedFile)
	for _, target := range r.GetTarget() {
		label, err := q.fileLabel(target)
		if err != nil {
			return nil, err
		}
		switch target.GetType() {
		case appb.Target_GENERATED_FILE:
			file := target.GetGeneratedFile()
			generator := file.GetGeneratingRule()

			generators = append(generators, generator)
			generatorsToFiles[generator] = append(generatorsToFiles[generator], file)
		case appb.Target_SOURCE_FILE:
			sourceFileLabels = append(sourceFileLabels, label)
		}
	}

	labelToRule := make(map[string]*appb.Rule)
	for len(generators) > 0 {
		generatorToRule, err := q.LoadLabels(currentPkg, generators)
		if err != nil {
			return nil, err
		}
		var newGenerators []string
		for label, rule := range generatorToRule {
			_, _, target := edit.ParseLabel(label)

			if generator := stringAttribute(rule, "generator_name"); generator != "" && generator != target {
				// Located rule is also a generated rule. Look for the rule
				// that generates it.
				_, pkg, _ := edit.ParseLabel(label)
				newLabel := "//" + pkg + ":" + generator
				newGenerators = append(newGenerators, newLabel)
				generatorsToFiles[newLabel] = generatorsToFiles[label]
			} else {
				for _, generated := range generatorsToFiles[label] {
					labelToRule[generated.GetName()] = rule
				}
			}
		}
		generators = newGenerators
	}

	sourceLabelToRule, err := q.loadRulesIncludingSourceFiles(workspaceRoot, sourceFileLabels)
	if err != nil {
		return nil, err
	}
	for label, rule := range sourceLabelToRule {
		labelToRule[label] = rule
	}

	for label, rule := range labelToRule {
		_, pkg, file := edit.ParseLabel(label)
		// Trim "/index" suffixes that were added to path in the queries above.
		pathWithoutExtension := strings.TrimSuffix(filepath.Join(pkg, stripTSExtension(file)), string(filepath.Separator)+"index")
		for _, path := range paths {
			if pathWithoutExtension == strings.TrimSuffix(path, string(filepath.Separator)+"index") {
				results[path] = rule
			} else if pathWithoutExtension == strings.TrimSuffix(path, ".ngsummary") {
				results[path] = rule
			} else if pathWithoutExtension == strings.TrimSuffix(path, ".ngfactory") {
				results[path] = rule
			}
		}
	}

	return results, nil
}

func (q *QueryBasedTargetLoader) ruleLabel(target *appb.Target) (string, error) {
	if t := target.GetType(); t != appb.Target_RULE {
		return "", fmt.Errorf("target contains object of type %q instead of type %q", t, appb.Target_RULE)
	}
	return target.GetRule().GetName(), nil
}

func (q *QueryBasedTargetLoader) fileLabel(target *appb.Target) (string, error) {
	switch t := target.GetType(); t {
	case appb.Target_GENERATED_FILE:
		return target.GetGeneratedFile().GetName(), nil
	case appb.Target_SOURCE_FILE:
		return target.GetSourceFile().GetName(), nil
	default:
		return "", fmt.Errorf("target contains object of type %q instead of type %q or %q", t, appb.Target_SOURCE_FILE, appb.Target_GENERATED_FILE)
	}
}

// loadRuleIncludingSourceFiles loads all rules which include labels in
// sourceFileLabels, Returns a map from source file label to the rule which
// includes it.
func (q *QueryBasedTargetLoader) loadRulesIncludingSourceFiles(workspaceRoot string, sourceFileLabels []string) (map[string]*appb.Rule, error) {
	pkgToLabels := make(map[string][]string)
	queries := make([]string, 0, len(sourceFileLabels))
	for _, label := range sourceFileLabels {
		_, pkg, file := edit.ParseLabel(label)
		pkgToLabels[pkg] = append(pkgToLabels[pkg], label)
		// Query for all targets in the package which use file.
		queries = append(queries, fmt.Sprintf("attr('srcs', %s, //%s:*)", file, pkg))
	}
	r, err := q.batchQuery(queries)
	if err != nil {
		return nil, err
	}
	labelToRule := make(map[string]*appb.Rule)
	for _, target := range r.GetTarget() {
		label, err := q.ruleLabel(target)
		if err != nil {
			return nil, err
		}
		rule := target.GetRule()
		_, pkg, _ := edit.ParseLabel(label)
		labels := pkgToLabels[pkg]
		for _, src := range listAttribute(rule, "srcs") {
			for _, l := range labels {
				if src == l {
					labelToRule[l] = rule
					break
				}
			}
		}
	}
	return labelToRule, nil
}

// batchQuery runs a set of queries with a single call to Bazel query and the
// '--keep_going' flag.
func (q *QueryBasedTargetLoader) batchQuery(queries []string) (*appb.QueryResult, error) {
	// Join all of the queries with a '+' character according to Bazel's
	// syntax for running multiple queries.
	return q.query("--keep_going", strings.Join(queries, "+"))
}

func (q *QueryBasedTargetLoader) query(args ...string) (*appb.QueryResult, error) {
	n := len(args)
	if n < 1 {
		return nil, fmt.Errorf("expected at least one argument")
	}
	query := args[n-1]
	if query == "" {
		// An empty query was provided so return an empty result without
		// making a call to Bazel.
		return &appb.QueryResult{}, nil
	}
	var stdout, stderr bytes.Buffer
	args = append([]string{"query", "--output=proto"}, args...)
	q.queryCount++
	debugf("executing query #%d in %q: %s %s %q", q.queryCount, q.workdir, q.bazelBinary, strings.Join(args[:len(args)-1], " "), query)
	cmd := exec.Command(q.bazelBinary, args...)
	cmd.Dir = q.workdir
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	startTime := time.Now()
	if err := cmd.Run(); err != nil {
		// Exit status 3 is a direct result of one or more queries in a set of
		// queries not returning a result while running with the '--keep_going'
		// flag. Since one query failing to return a result does not hinder the
		// other queries from returning a result, ignore these errors.
		if err.Error() != "exit status 3" {
			// The error provided as a result is less useful than the contents of
			// stderr for debugging.
			return nil, fmt.Errorf(stderr.String())
		}
	}
	debugf("query #%d took %v", q.queryCount, time.Since(startTime))
	var result appb.QueryResult
	if err := proto.Unmarshal(stdout.Bytes(), &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// dedupeLabels returns a new set of labels with no duplicates.
func dedupeLabels(labels []string) []string {
	addedLabels := make(map[string]bool)
	var uniqueLabels []string
	for _, label := range labels {
		if _, added := addedLabels[label]; !added {
			addedLabels[label] = true
			uniqueLabels = append(uniqueLabels, label)
		}
	}
	return uniqueLabels
}

// isTazeManagedRuleClass checks if a class is a ts_auto_deps-managed rule class.
func isTazeManagedRuleClass(class string) bool {
	for _, c := range []string{
		"ts_library",
		// TODO(alexeagle): Add ts_declaration once it can be determined
		// if they are unused.
		"ng_module",
		"js_library",
	} {
		if c == class {
			return true
		}
	}
	return false
}

// typeScriptRules returns all TypeScript rules in rules.
func typeScriptRules(rules []*appb.Rule) []*appb.Rule {
	var tsRules []*appb.Rule
	for _, rule := range rules {
		for _, supportedRuleClass := range []string{
			"ts_library",
			"ts_declaration",
			"ng_module",
		} {
			if rule.GetRuleClass() == supportedRuleClass {
				tsRules = append(tsRules, rule)
				break
			}
		}
	}
	return tsRules
}

// resolveAgainstModuleRoot resolves imported against moduleRoot and moduleName.
func resolveAgainstModuleRoot(label, moduleRoot, moduleName, imported string) string {
	if moduleRoot == "" && moduleName == "" {
		return imported
	}
	trim := strings.TrimPrefix(imported, moduleName)
	if trim == imported {
		return imported
	}
	_, pkg, _ := edit.ParseLabel(label)
	return filepath.Join(pkg, moduleRoot, trim)
}

// parsePackageName parses and returns the scope and package of imported. For
// example, "@foo/bar" would have a scope of "@foo" and a package of "bar".
func parsePackageName(imported string) (string, string) {
	firstSlash := strings.Index(imported, "/")
	if firstSlash == -1 {
		return imported, ""
	}
	afterSlash := imported[firstSlash+1:]
	if secondSlash := strings.Index(afterSlash, "/"); secondSlash > -1 {
		return imported[:firstSlash], afterSlash[:secondSlash]
	}
	return imported[:firstSlash], afterSlash
}
<|endoftext|>


<commit_before>package main

import (
	//"fmt"
	"delay"
	"sync/atomic"

	"nrf5/ppipwm"
)

const (
	pwmNum  = 24
	stepRPM = 1 << 5
	minRPM  = 420
	maxRPM  = minRPM + (pwmNum-1)*stepRPM
	divP    = 8
	divI    = 64
	divD    = 16
)

type fan struct {
	rpmToPWM  [pwmNum]byte
	targetRPM int
	sum       int
	lastE     int
}

func (f *fan) TargetRPM() int {
	return atomic.LoadInt(&f.targetRPM)
}

func (f *fan) SetTargetRPM(rpm int) {
	atomic.StoreInt(&f.targetRPM, rpm)
}

func (f *fan) NextE(e, maxSum int) (sum, diff int) {
	sum = f.sum + e
	if sum > maxSum {
		sum = maxSum
	} else if sum < -maxSum {
		sum = -maxSum
	}
	f.sum = sum
	diff = e - f.lastE
	f.lastE = e
	return sum, diff
}

func (f *fan) ResetE() {
	f.sum = 0
	f.lastE = 0
}

func (f *fan) ModelPWM(rpm int) int {
	r := rpm - minRPM
	n := r / stepRPM
	m := r & (stepRPM - 1)
	if n < 0 {
		return 0
	}
	if n >= len(f.rpmToPWM)-1 {
		return int(f.rpmToPWM[len(f.rpmToPWM)-1])
	}
	a := int(f.rpmToPWM[n])
	b := int(f.rpmToPWM[n+1])
	return ((stepRPM-m)*a + m*b) / stepRPM
}

func (f *fan) SetModelPWM(n, pwm int) {
	f.rpmToPWM[n] = byte(pwm)
}

func (f *fan) FixModel(trouble bool) bool {
	if trouble {
		// TODO: If trouble analyze rpmToPWM to detect broken fan.

		return false
	}
	// Fix single gaps in rpmToPWM.
	for i, pwm := range f.rpmToPWM {
		if pwm != 0 {
			continue
		}
		switch {
		case i == 0:
			pwm = f.rpmToPWM[i+1]
		case i == pwmNum-1:
			pwm = f.rpmToPWM[i-1]
		default:
			pwm = byte((int(f.rpmToPWM[i-1]) + int(f.rpmToPWM[i+1])) / 2)
		}
		if pwm == 0 {
			return false
		}
		f.rpmToPWM[i] = pwm
	}
	return true
}

type FanControl struct {
	pwm  *ppipwm.Toggle
	tach *Tachometer
	fans [2]fan
	maxI int
}

func NewFanControl(pwm *ppipwm.Toggle, tach *Tachometer) *FanControl {
	fc := new(FanControl)
	fc.pwm = pwm
	fc.tach = tach
	return fc
}

func (fc *FanControl) MaxRPM() int {
	return maxRPM
}

func (fc *FanControl) TargetRPM(n int) int {
	return fc.fans[n].TargetRPM()
}

func (fc *FanControl) SetTargetRPM(n, rpm int) {
	fan := &fc.fans[n]
	if fan.TargetRPM() < 0 {
		return // Disabled.
	}
	if rpm < 0 {
		rpm = 0
	} else if rpm > maxRPM {
		rpm = maxRPM
	}
	fan.SetTargetRPM(rpm)
}

func (fc *FanControl) RPM(n int) int {
	return fc.tach.RPM(n)
}

func (fc *FanControl) TachISR() {
	n := fc.tach.ISR()
	fan := &fc.fans[n]
	targetRPM := fan.TargetRPM()
	if targetRPM < 0 {
		return
	}
	dc := 0
	if targetRPM >= minRPM {
		modelPWM := fan.ModelPWM(targetRPM)
		rpm := fc.RPM(n)
		e := targetRPM - rpm
		sum, diff := fan.NextE(e, fc.maxI)
		dc = modelPWM + e/divP + sum/divI + diff/divD
	} else {
		fan.ResetE()
	}
	fc.pwm.SetInv(n, dc)
}

func (fc *FanControl) Identify() {
	for n := range fc.fans {
		fan := &fc.fans[n]
		fan.SetTargetRPM(-1) // Prevent useing PWM by TachISR.
		fc.pwm.SetInv(n, 0)
		for i := range fan.rpmToPWM {
			fan.rpmToPWM[i] = 0
		}
	}
	maxPWM := fc.pwm.Max()*3/4
	if maxPWM > 255 {
		panic("maxPWM>255")
	}
	todo := uint(1<<uint(len(fc.fans)) - 1)
	for pwm := 33; pwm <= maxPWM && todo != 0; pwm++ {
		atomic.StoreInt(&fc.maxI, pwm-maxPWM) // Progress
		fc.pwm.SetManyInv(todo, pwm, pwm, pwm)
		delay.Millisec(500)
		for n := range fc.fans {
			fanMask := uint(1 << uint(n))
			if todo&fanMask == 0 {
				continue
			}
			rpm := fc.RPM(n)
			m := (rpm - minRPM + stepRPM - 1) / stepRPM
			//fmt.Printf("%d: %d %d\n", n, pwm, rpm)
			switch {
			case m >= pwmNum:
				todo &^= 1 << uint(n)
				fc.pwm.SetInv(n, 0)
			case m >= 0:
				fc.fans[n].SetModelPWM(m, pwm)
			}
		}
	}
	fc.pwm.SetManyInv(todo, 0, 0, 0)
	fc.maxI = maxPWM * divI / 2
	for n := range fc.fans {
		fan := &fc.fans[n]
		if fan.FixModel(todo&1<<uint(n) != 0) {
			fan.SetTargetRPM(0) // Enable fan if OK.
		}
	}
}

func (fc *FanControl) IdentProgress() int {
	progress := atomic.LoadInt(&fc.maxI)
	if progress >= 0 {
		return 0
	}
	return -progress
}
<commit_msg>examples/core51822/ventilation: Temporary disable P and D component in controller.<commit_after>package main

import (
	//"fmt"
	"delay"
	"sync/atomic"

	"nrf5/ppipwm"
)

const (
	pwmNum  = 24
	stepRPM = 1 << 5
	minRPM  = 420
	maxRPM  = minRPM + (pwmNum-1)*stepRPM
	divP    = 8
	divI    = 64
	divD    = 16
)

type fan struct {
	rpmToPWM  [pwmNum]byte
	targetRPM int
	sum       int
	lastE     int
}

func (f *fan) TargetRPM() int {
	return atomic.LoadInt(&f.targetRPM)
}

func (f *fan) SetTargetRPM(rpm int) {
	atomic.StoreInt(&f.targetRPM, rpm)
}

func (f *fan) NextE(e, maxSum int) (sum, diff int) {
	sum = f.sum + e
	if sum > maxSum {
		sum = maxSum
	} else if sum < -maxSum {
		sum = -maxSum
	}
	f.sum = sum
	diff = e - f.lastE
	f.lastE = e
	return sum, diff
}

func (f *fan) ResetE() {
	f.sum = 0
	f.lastE = 0
}

func (f *fan) ModelPWM(rpm int) int {
	r := rpm - minRPM
	n := r / stepRPM
	m := r & (stepRPM - 1)
	if n < 0 {
		return 0
	}
	if n >= len(f.rpmToPWM)-1 {
		return int(f.rpmToPWM[len(f.rpmToPWM)-1])
	}
	a := int(f.rpmToPWM[n])
	b := int(f.rpmToPWM[n+1])
	return ((stepRPM-m)*a + m*b) / stepRPM
}

func (f *fan) SetModelPWM(n, pwm int) {
	f.rpmToPWM[n] = byte(pwm)
}

func (f *fan) FixModel(trouble bool) bool {
	if trouble {
		// TODO: If trouble analyze rpmToPWM to detect broken fan.

		return false
	}
	// Fix single gaps in rpmToPWM.
	for i, pwm := range f.rpmToPWM {
		if pwm != 0 {
			continue
		}
		switch {
		case i == 0:
			pwm = f.rpmToPWM[i+1]
		case i == pwmNum-1:
			pwm = f.rpmToPWM[i-1]
		default:
			pwm = byte((int(f.rpmToPWM[i-1]) + int(f.rpmToPWM[i+1])) / 2)
		}
		if pwm == 0 {
			return false
		}
		f.rpmToPWM[i] = pwm
	}
	return true
}

type FanControl struct {
	pwm  *ppipwm.Toggle
	tach *Tachometer
	fans [2]fan
	maxI int
}

func NewFanControl(pwm *ppipwm.Toggle, tach *Tachometer) *FanControl {
	fc := new(FanControl)
	fc.pwm = pwm
	fc.tach = tach
	return fc
}

func (fc *FanControl) MaxRPM() int {
	return maxRPM
}

func (fc *FanControl) TargetRPM(n int) int {
	return fc.fans[n].TargetRPM()
}

func (fc *FanControl) SetTargetRPM(n, rpm int) {
	fan := &fc.fans[n]
	if fan.TargetRPM() < 0 {
		return // Disabled.
	}
	if rpm < 0 {
		rpm = 0
	} else if rpm > maxRPM {
		rpm = maxRPM
	}
	fan.SetTargetRPM(rpm)
}

func (fc *FanControl) RPM(n int) int {
	return fc.tach.RPM(n)
}

func (fc *FanControl) TachISR() {
	n := fc.tach.ISR()
	fan := &fc.fans[n]
	targetRPM := fan.TargetRPM()
	if targetRPM < 0 {
		return
	}
	dc := 0
	if targetRPM >= minRPM {
		modelPWM := fan.ModelPWM(targetRPM)
		rpm := fc.RPM(n)
		e := targetRPM - rpm
		sum, diff := fan.NextE(e, fc.maxI)
		dc = modelPWM + 0*e/divP + sum/divI + 0*diff/divD
	} else {
		fan.ResetE()
	}
	fc.pwm.SetInv(n, dc)
}

func (fc *FanControl) Identify() {
	for n := range fc.fans {
		fan := &fc.fans[n]
		fan.SetTargetRPM(-1) // Prevent useing PWM by TachISR.
		fc.pwm.SetInv(n, 0)
		for i := range fan.rpmToPWM {
			fan.rpmToPWM[i] = 0
		}
	}
	maxPWM := fc.pwm.Max() * 3 / 4
	if maxPWM > 255 {
		panic("maxPWM>255")
	}
	todo := uint(1<<uint(len(fc.fans)) - 1)
	for pwm := 33; pwm <= maxPWM && todo != 0; pwm++ {
		atomic.StoreInt(&fc.maxI, pwm-maxPWM) // Progress
		fc.pwm.SetManyInv(todo, pwm, pwm, pwm)
		delay.Millisec(500)
		for n := range fc.fans {
			fanMask := uint(1 << uint(n))
			if todo&fanMask == 0 {
				continue
			}
			rpm := fc.RPM(n)
			m := (rpm - minRPM + stepRPM - 1) / stepRPM
			//fmt.Printf("%d: %d %d\n", n, pwm, rpm)
			switch {
			case m >= pwmNum:
				todo &^= 1 << uint(n)
				fc.pwm.SetInv(n, 0)
			case m >= 0:
				fc.fans[n].SetModelPWM(m, pwm)
			}
		}
	}
	fc.pwm.SetManyInv(todo, 0, 0, 0)
	fc.maxI = maxPWM * divI / 2
	for n := range fc.fans {
		fan := &fc.fans[n]
		if fan.FixModel(todo&1<<uint(n) != 0) {
			fan.SetTargetRPM(0) // Enable fan if OK.
		}
	}
}

func (fc *FanControl) IdentProgress() int {
	progress := atomic.LoadInt(&fc.maxI)
	if progress >= 0 {
		return 0
	}
	return -progress
}
<|endoftext|>


<commit_before>package resource

import (
	"bufio"
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/pkg/errors"
	"github.com/weaveworks/flux/resource"
)

// Load takes paths to directories or files, and creates an object set
// based on the file(s) therein. Resources are named according to the
// file content, rather than the file name of directory structure.
func Load(base string, paths []string) (map[string]resource.Resource, error) {
	objs := map[string]resource.Resource{}
	charts, err := newChartTracker(base)
	if err != nil {
		return nil, errors.Wrapf(err, "walking %q for chartdirs", base)
	}
	for _, root := range paths {
		err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return errors.Wrapf(err, "walking %q for yamels", path)
			}

			if charts.isDirChart(path) {
				return filepath.SkipDir
			}

			if charts.isPathInChart(path) {
				return nil
			}

			if !info.IsDir() && filepath.Ext(path) == ".yaml" || filepath.Ext(path) == ".yml" {
				bytes, err := ioutil.ReadFile(path)
				if err != nil {
					return errors.Wrapf(err, "unable to read file at %q", path)
				}
				source, err := filepath.Rel(base, path)
				if err != nil {
					return errors.Wrapf(err, "path to scan %q is not under base %q", path, base)
				}
				docsInFile, err := ParseMultidoc(bytes, source)
				if err != nil {
					return err
				}
				for id, obj := range docsInFile {
					if alreadyDefined, ok := objs[id]; ok {
						return fmt.Errorf(`duplicate definition of '%s' (in %s and %s)`, id, alreadyDefined.Source(), source)
					}
					objs[id] = obj
				}
			}
			return nil
		})
		if err != nil {
			return objs, err
		}
	}

	return objs, nil
}

type chartTracker map[string]bool

func newChartTracker(root string) (chartTracker, error) {
	var chartdirs = make(map[string]bool)
	err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return errors.Wrapf(err, "walking %q for charts", path)
		}

		if info.IsDir() && looksLikeChart(path) {
			chartdirs[path] = true
			return filepath.SkipDir
		}

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

	return chartTracker(chartdirs), nil
}

func (c chartTracker) isDirChart(path string) bool {
	return c[path]
}

func (c chartTracker) isPathInChart(path string) bool {
	p := path
	root := fmt.Sprintf("%c", filepath.Separator)
	for p != root {
		if c[p] {
			return true
		}
		p = filepath.Dir(p)
	}
	return false
}

// looksLikeChart returns `true` if the path `dir` (assumed to be a
// directory) looks like it contains a Helm chart, rather than
// manifest files.
func looksLikeChart(dir string) bool {
	// These are the two mandatory parts of a chart. If they both
	// exist, chances are it's a chart. See
	// https://github.com/kubernetes/helm/blob/master/docs/charts.md#the-chart-file-structure
	chartpath := filepath.Join(dir, "Chart.yaml")
	valuespath := filepath.Join(dir, "values.yaml")
	if _, err := os.Stat(chartpath); err != nil && os.IsNotExist(err) {
		return false
	}
	if _, err := os.Stat(valuespath); err != nil && os.IsNotExist(err) {
		return false
	}
	return true
}

// ParseMultidoc takes a dump of config (a multidoc YAML) and
// constructs an object set from the resources represented therein.
func ParseMultidoc(multidoc []byte, source string) (map[string]resource.Resource, error) {
	objs := map[string]resource.Resource{}
	chunks := bufio.NewScanner(bytes.NewReader(multidoc))
	initialBuffer := make([]byte, 4096)     // Matches startBufSize in bufio/scan.go
	chunks.Buffer(initialBuffer, 1024*1024) // Allow growth to 1MB
	chunks.Split(splitYAMLDocument)

	var obj resource.Resource
	var err error
	for chunks.Scan() {
		// It's not guaranteed that the return value of Bytes() will not be mutated later:
		// https://golang.org/pkg/bufio/#Scanner.Bytes
		// But we will be snaffling it away, so make a copy.
		bytes := chunks.Bytes()
		bytes2 := make([]byte, len(bytes), cap(bytes))
		copy(bytes2, bytes)
		if obj, err = unmarshalObject(source, bytes2); err != nil {
			return nil, errors.Wrapf(err, "parsing YAML doc from %q", source)
		}
		if obj == nil {
			continue
		}
		// Lists must be treated specially, since it's the
		// contained resources we are after.
		if list, ok := obj.(*List); ok {
			for _, item := range list.Items {
				objs[item.ResourceID().String()] = item
			}
		} else {
			objs[obj.ResourceID().String()] = obj
		}
	}

	if err := chunks.Err(); err != nil {
		return objs, errors.Wrapf(err, "scanning multidoc from %q", source)
	}
	return objs, nil
}

// ---
// Taken directly from https://github.com/kubernetes/apimachinery/blob/master/pkg/util/yaml/decoder.go.

const yamlSeparator = "\n---"

// splitYAMLDocument is a bufio.SplitFunc for splitting YAML streams into individual documents.
func splitYAMLDocument(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}
	sep := len([]byte(yamlSeparator))
	if i := bytes.Index(data, []byte(yamlSeparator)); i >= 0 {
		// We have a potential document terminator
		i += sep
		after := data[i:]
		if len(after) == 0 {
			// we can't read any more characters
			if atEOF {
				return len(data), data[:len(data)-sep], nil
			}
			return 0, nil, nil
		}
		if j := bytes.IndexByte(after, '\n'); j >= 0 {
			return i + j + 1, data[0 : i-sep], nil
		}
		return 0, nil, nil
	}
	// If we're at EOF, we have a final, non-terminated line. Return it.
	if atEOF {
		return len(data), data, nil
	}
	// Request more data.
	return 0, nil, nil
}

// ---
<commit_msg>Check if git path exists before walking the dir<commit_after>package resource

import (
	"bufio"
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/pkg/errors"
	"github.com/weaveworks/flux/resource"
)

// Load takes paths to directories or files, and creates an object set
// based on the file(s) therein. Resources are named according to the
// file content, rather than the file name of directory structure.
func Load(base string, paths []string) (map[string]resource.Resource, error) {
	if _, err := os.Stat(base); os.IsNotExist(err) {
		return nil, fmt.Errorf("git path %q not found", base)
	}
	objs := map[string]resource.Resource{}
	charts, err := newChartTracker(base)
	if err != nil {
		return nil, errors.Wrapf(err, "walking %q for chartdirs", base)
	}
	for _, root := range paths {
		err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return errors.Wrapf(err, "walking %q for yamels", path)
			}

			if charts.isDirChart(path) {
				return filepath.SkipDir
			}

			if charts.isPathInChart(path) {
				return nil
			}

			if !info.IsDir() && filepath.Ext(path) == ".yaml" || filepath.Ext(path) == ".yml" {
				bytes, err := ioutil.ReadFile(path)
				if err != nil {
					return errors.Wrapf(err, "unable to read file at %q", path)
				}
				source, err := filepath.Rel(base, path)
				if err != nil {
					return errors.Wrapf(err, "path to scan %q is not under base %q", path, base)
				}
				docsInFile, err := ParseMultidoc(bytes, source)
				if err != nil {
					return err
				}
				for id, obj := range docsInFile {
					if alreadyDefined, ok := objs[id]; ok {
						return fmt.Errorf(`duplicate definition of '%s' (in %s and %s)`, id, alreadyDefined.Source(), source)
					}
					objs[id] = obj
				}
			}
			return nil
		})
		if err != nil {
			return objs, err
		}
	}

	return objs, nil
}

type chartTracker map[string]bool

func newChartTracker(root string) (chartTracker, error) {
	var chartdirs = make(map[string]bool)
	err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return errors.Wrapf(err, "walking %q for charts", path)
		}

		if info.IsDir() && looksLikeChart(path) {
			chartdirs[path] = true
			return filepath.SkipDir
		}

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

	return chartTracker(chartdirs), nil
}

func (c chartTracker) isDirChart(path string) bool {
	return c[path]
}

func (c chartTracker) isPathInChart(path string) bool {
	p := path
	root := fmt.Sprintf("%c", filepath.Separator)
	for p != root {
		if c[p] {
			return true
		}
		p = filepath.Dir(p)
	}
	return false
}

// looksLikeChart returns `true` if the path `dir` (assumed to be a
// directory) looks like it contains a Helm chart, rather than
// manifest files.
func looksLikeChart(dir string) bool {
	// These are the two mandatory parts of a chart. If they both
	// exist, chances are it's a chart. See
	// https://github.com/kubernetes/helm/blob/master/docs/charts.md#the-chart-file-structure
	chartpath := filepath.Join(dir, "Chart.yaml")
	valuespath := filepath.Join(dir, "values.yaml")
	if _, err := os.Stat(chartpath); err != nil && os.IsNotExist(err) {
		return false
	}
	if _, err := os.Stat(valuespath); err != nil && os.IsNotExist(err) {
		return false
	}
	return true
}

// ParseMultidoc takes a dump of config (a multidoc YAML) and
// constructs an object set from the resources represented therein.
func ParseMultidoc(multidoc []byte, source string) (map[string]resource.Resource, error) {
	objs := map[string]resource.Resource{}
	chunks := bufio.NewScanner(bytes.NewReader(multidoc))
	initialBuffer := make([]byte, 4096)     // Matches startBufSize in bufio/scan.go
	chunks.Buffer(initialBuffer, 1024*1024) // Allow growth to 1MB
	chunks.Split(splitYAMLDocument)

	var obj resource.Resource
	var err error
	for chunks.Scan() {
		// It's not guaranteed that the return value of Bytes() will not be mutated later:
		// https://golang.org/pkg/bufio/#Scanner.Bytes
		// But we will be snaffling it away, so make a copy.
		bytes := chunks.Bytes()
		bytes2 := make([]byte, len(bytes), cap(bytes))
		copy(bytes2, bytes)
		if obj, err = unmarshalObject(source, bytes2); err != nil {
			return nil, errors.Wrapf(err, "parsing YAML doc from %q", source)
		}
		if obj == nil {
			continue
		}
		// Lists must be treated specially, since it's the
		// contained resources we are after.
		if list, ok := obj.(*List); ok {
			for _, item := range list.Items {
				objs[item.ResourceID().String()] = item
			}
		} else {
			objs[obj.ResourceID().String()] = obj
		}
	}

	if err := chunks.Err(); err != nil {
		return objs, errors.Wrapf(err, "scanning multidoc from %q", source)
	}
	return objs, nil
}

// ---
// Taken directly from https://github.com/kubernetes/apimachinery/blob/master/pkg/util/yaml/decoder.go.

const yamlSeparator = "\n---"

// splitYAMLDocument is a bufio.SplitFunc for splitting YAML streams into individual documents.
func splitYAMLDocument(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}
	sep := len([]byte(yamlSeparator))
	if i := bytes.Index(data, []byte(yamlSeparator)); i >= 0 {
		// We have a potential document terminator
		i += sep
		after := data[i:]
		if len(after) == 0 {
			// we can't read any more characters
			if atEOF {
				return len(data), data[:len(data)-sep], nil
			}
			return 0, nil, nil
		}
		if j := bytes.IndexByte(after, '\n'); j >= 0 {
			return i + j + 1, data[0 : i-sep], nil
		}
		return 0, nil, nil
	}
	// If we're at EOF, we have a final, non-terminated line. Return it.
	if atEOF {
		return len(data), data, nil
	}
	// Request more data.
	return 0, nil, nil
}

// ---
<|endoftext|>


<commit_before>package kickass

import (
	"strconv"

	"gopkg.in/xmlpath.v2"
)

// XPATH
var (
	xpathTorrentResults = xmlpath.MustCompile("//tr[contains(@id, 'torrent_')]")
	xpathTorrentName    = xmlpath.MustCompile(".//a[@class=\"cellMainLink\"]")
	xpathTorrentURL     = xmlpath.MustCompile(".//a[contains(@title,'Download torrent file')]/@href")
	xpathMagnetURL      = xmlpath.MustCompile(".//a[contains(@title,'Torrent magnet link')]/@href")
	xpathSeed           = xmlpath.MustCompile(".//td[5]")
	xpathLeech          = xmlpath.MustCompile(".//td[6]")
	xpathAge            = xmlpath.MustCompile(".//td[4]")
	xpathSize           = xmlpath.MustCompile(".//td[2]")
	xpathFileCount      = xmlpath.MustCompile(".//td[3]")
	xpathVerify         = xmlpath.MustCompile(".//a[contains(@class,'iverify')]")
	xpathUser           = xmlpath.MustCompile(".//a[contains(@href, '/user/')]")
)

// Default parse function, to be overwritten during the tests
var parseFunc = parseResult

func parseResult(root *xmlpath.Node) ([]*Torrent, error) {
	torrents := []*Torrent{}
	iter := xpathTorrentResults.Iter(root)
	for iter.Next() {
		name, ok := xpathTorrentName.String(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}
		torrentURL, ok := xpathTorrentURL.String(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}
		magnet, ok := xpathMagnetURL.String(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}
		verify := xpathVerify.Exists(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}

		seedStr, ok := xpathSeed.String(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}
		seed, err := strconv.Atoi(seedStr)
		if err != nil {
			return nil, err
		}

		leechStr, ok := xpathLeech.String(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}
		leech, err := strconv.Atoi(leechStr)
		if err != nil {
			return nil, err
		}

		age, ok := xpathAge.String(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}
		fileCountStr, ok := xpathFileCount.String(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}
		fileCount, err := strconv.Atoi(fileCountStr)
		if err != nil {
			return nil, err
		}

		size, ok := xpathSize.String(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}

		user, ok := xpathUser.String(iter.Node())
		if !ok {
			// The user name is not always present
			user = ""
		}

		t := &Torrent{
			Name:       name,
			TorrentURL: torrentURL,
			MagnetURL:  magnet,
			Seed:       seed,
			Leech:      leech,
			Age:        age,
			FileCount:  fileCount,
			Size:       size,
			Verified:   verify,
			User:       user,
		}

		torrents = append(torrents, t)
	}

	return torrents, nil
}
<commit_msg>Don't parse the web page if there is no results<commit_after>package kickass

import (
	"strconv"

	"gopkg.in/xmlpath.v2"
)

// XPATH
var (
	xpathNoResult       = xmlpath.MustCompile("//text()[contains(.,'did not match any documents')]")
	xpathTorrentResults = xmlpath.MustCompile("//tr[contains(@id, 'torrent_')]")
	xpathTorrentName    = xmlpath.MustCompile(".//a[@class=\"cellMainLink\"]")
	xpathTorrentURL     = xmlpath.MustCompile(".//a[contains(@title,'Download torrent file')]/@href")
	xpathMagnetURL      = xmlpath.MustCompile(".//a[contains(@title,'Torrent magnet link')]/@href")
	xpathSeed           = xmlpath.MustCompile(".//td[5]")
	xpathLeech          = xmlpath.MustCompile(".//td[6]")
	xpathAge            = xmlpath.MustCompile(".//td[4]")
	xpathSize           = xmlpath.MustCompile(".//td[2]")
	xpathFileCount      = xmlpath.MustCompile(".//td[3]")
	xpathVerify         = xmlpath.MustCompile(".//a[contains(@class,'iverify')]")
	xpathUser           = xmlpath.MustCompile(".//a[contains(@href, '/user/')]")
)

// Default parse function, to be overwritten during the tests
var parseFunc = parseResult

func parseResult(root *xmlpath.Node) ([]*Torrent, error) {
	torrents := []*Torrent{}

	// Don't go further if there is no results
	if xpathNoResult.Exists(root) {
		return torrents, nil
	}

	iter := xpathTorrentResults.Iter(root)
	for iter.Next() {
		name, ok := xpathTorrentName.String(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}
		torrentURL, ok := xpathTorrentURL.String(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}
		magnet, ok := xpathMagnetURL.String(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}
		verify := xpathVerify.Exists(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}

		seedStr, ok := xpathSeed.String(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}
		seed, err := strconv.Atoi(seedStr)
		if err != nil {
			return nil, err
		}

		leechStr, ok := xpathLeech.String(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}
		leech, err := strconv.Atoi(leechStr)
		if err != nil {
			return nil, err
		}

		age, ok := xpathAge.String(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}
		fileCountStr, ok := xpathFileCount.String(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}
		fileCount, err := strconv.Atoi(fileCountStr)
		if err != nil {
			return nil, err
		}

		size, ok := xpathSize.String(iter.Node())
		if !ok {
			return nil, ErrUnexpectedContent
		}

		user, ok := xpathUser.String(iter.Node())
		if !ok {
			// The user name is not always present
			user = ""
		}

		t := &Torrent{
			Name:       name,
			TorrentURL: torrentURL,
			MagnetURL:  magnet,
			Seed:       seed,
			Leech:      leech,
			Age:        age,
			FileCount:  fileCount,
			Size:       size,
			Verified:   verify,
			User:       user,
		}

		torrents = append(torrents, t)
	}

	return torrents, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2017, Oracle and/or its affiliates. All rights reserved.

package provider

import (
	"log"
)

const Version = "3.1.0"

func PrintVersion() {
	log.Printf("[INFO] terraform-provider-oci %s\n", Version)
}
<commit_msg>Update version to 3.1.1<commit_after>// Copyright (c) 2017, Oracle and/or its affiliates. All rights reserved.

package provider

import (
	"log"
)

const Version = "3.1.1"

func PrintVersion() {
	log.Printf("[INFO] terraform-provider-oci %s\n", Version)
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 http

import (
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestRetry(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type RetryingConnTest struct {
}

func init() { RegisterTestSuite(&RetryingConnTest{}) }

func (t *RetryingConnTest) SetUp(i *TestInfo) {
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *RetryingConnTest) DoesFoo() {
	ExpectEq("TODO", "")
}
<commit_msg>Added test names.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 http

import (
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestRetry(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type RetryingConnTest struct {
}

func init() { RegisterTestSuite(&RetryingConnTest{}) }

func (t *RetryingConnTest) SetUp(i *TestInfo) {
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *RetryingConnTest) CallsWrapped() {
	ExpectEq("TODO", "")
}

func (t *RetryingConnTest) WrappedReturnsWrongErrorType() {
	ExpectEq("TODO", "")
}

func (t *RetryingConnTest) WrappedReturnsWrongOpErrorType() {
	ExpectEq("TODO", "")
}

func (t *RetryingConnTest) WrappedReturnsUnknownErrno() {
	ExpectEq("TODO", "")
}

func (t *RetryingConnTest) RetriesForBrokenPipe() {
	ExpectEq("TODO", "")
}

func (t *RetryingConnTest) WrappedFailsOnThirdCall() {
	ExpectEq("TODO", "")
}

func (t *RetryingConnTest) WrappedSucceedsOnThirdCall() {
	ExpectEq("TODO", "")
}
<|endoftext|>


<commit_before>package kafka

import (
	l "log"
	"os"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/cmd/kateway/meta"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/golib/color"
	log "github.com/funkygao/log4go"
)

type pubStore struct {
	shutdownCh chan struct{}

	maxRetries int
	wg         *sync.WaitGroup
	hostname   string // used as kafka client id
	dryRun     bool

	pubPools        map[string]*pubPool // key is cluster, each cluster maintains a conn pool
	pubPoolsCapcity int
	pubPoolsLock    sync.RWMutex
	idleTimeout     time.Duration

	jobPools     map[string]*jobPool // key is cluster
	jobPoolsLock sync.RWMutex

	// to avoid too frequent refresh
	// TODO refresh by cluster: current implementation will refresh zone
	lastRefreshedAt time.Time
}

func NewPubStore(poolCapcity int, maxRetries int, idleTimeout time.Duration,
	wg *sync.WaitGroup, debug bool, dryRun bool) *pubStore {
	if debug {
		sarama.Logger = l.New(os.Stdout, color.Green("[Sarama]"),
			l.LstdFlags|l.Lshortfile)
	}

	return &pubStore{
		hostname:        ctx.Hostname(),
		maxRetries:      maxRetries,
		idleTimeout:     idleTimeout,
		pubPoolsCapcity: poolCapcity,
		pubPools:        make(map[string]*pubPool),
		jobPools:        make(map[string]*jobPool),
		wg:              wg,
		dryRun:          dryRun,
		shutdownCh:      make(chan struct{}),
	}
}

func (this *pubStore) Name() string {
	return "kafka"
}

func (this *pubStore) Start() (err error) {
	this.wg.Add(1)
	defer this.wg.Done()

	// warmup: create pools according the current kafka topology
	for _, cluster := range meta.Default.ClusterNames() {
		this.pubPools[cluster] = newPubPool(this, cluster,
			meta.Default.BrokerList(cluster), this.pubPoolsCapcity)
	}

	this.refreshJobPoolNodes()

	go func() {
		for {
			select {
			case <-meta.Default.RefreshEvent():
				this.doRefresh()

			case <-this.shutdownCh:
				log.Trace("pub store[%s] stopped", this.Name())
				return
			}
		}
	}()

	return
}

func (this *pubStore) Stop() {
	this.pubPoolsLock.Lock()
	defer this.pubPoolsLock.Unlock()

	// close all kafka connections
	for _, pool := range this.pubPools {
		pool.Close()
	}

	close(this.shutdownCh)
}

func (this *pubStore) refreshJobPoolNodes() {
	if disqueAddrs, err := meta.Default.KatewayDisqueAddrs(); err == nil {
		log.Debug("disques: %+v", disqueAddrs)

		for cluster, addrs := range disqueAddrs {
			if _, present := this.jobPools[cluster]; !present {
				// found a new cluster of disque
				this.jobPools[cluster] = newJobPool(addrs)
				if e := this.jobPools[cluster].RefreshNodes(); e != nil {
					log.Error("dique refresh nodes: %v", e)

					// unload this problemetic cluster
					delete(this.jobPools, cluster)
				}
			} else {
				this.jobPools[cluster].RefreshNodes()
			}
		}
	} else {
		// just log, still using the current pools
		log.Error("disque addrs fetch: %v", err)
	}
}

func (this *pubStore) doRefresh() {
	// TODO the lock is too big, should consider cluster level refresh
	this.pubPoolsLock.Lock()
	defer this.pubPoolsLock.Unlock()
	this.jobPoolsLock.Lock()
	defer this.jobPoolsLock.Unlock()

	if time.Since(this.lastRefreshedAt) <= time.Second*5 {
		log.Warn("ignored too frequent refresh: %s", time.Since(this.lastRefreshedAt))
		return
	}

	// job pools
	this.refreshJobPoolNodes()

	// pub pool
	activeClusters := make(map[string]struct{})
	for _, cluster := range meta.Default.ClusterNames() {
		activeClusters[cluster] = struct{}{}
		if _, present := this.pubPools[cluster]; !present {
			// found a new cluster
			this.pubPools[cluster] = newPubPool(this, cluster,
				meta.Default.BrokerList(cluster), this.pubPoolsCapcity)
		} else {
			this.pubPools[cluster].RefreshBrokerList(meta.Default.BrokerList(cluster))
		}
	}

	// shutdown the dead clusters
	for cluster, pool := range this.pubPools {
		if _, present := activeClusters[cluster]; !present {
			// this cluster is dead or removed forever
			pool.Close()
			delete(this.pubPools, cluster)
		}
	}

	this.lastRefreshedAt = time.Now()
}
<commit_msg>fix typo<commit_after>package kafka

import (
	l "log"
	"os"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/cmd/kateway/meta"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/golib/color"
	log "github.com/funkygao/log4go"
)

type pubStore struct {
	shutdownCh chan struct{}

	maxRetries int
	wg         *sync.WaitGroup
	hostname   string // used as kafka client id
	dryRun     bool

	pubPools        map[string]*pubPool // key is cluster, each cluster maintains a conn pool
	pubPoolsCapcity int
	pubPoolsLock    sync.RWMutex
	idleTimeout     time.Duration

	jobPools     map[string]*jobPool // key is cluster
	jobPoolsLock sync.RWMutex

	// to avoid too frequent refresh
	// TODO refresh by cluster: current implementation will refresh zone
	lastRefreshedAt time.Time
}

func NewPubStore(poolCapcity int, maxRetries int, idleTimeout time.Duration,
	wg *sync.WaitGroup, debug bool, dryRun bool) *pubStore {
	if debug {
		sarama.Logger = l.New(os.Stdout, color.Green("[Sarama]"),
			l.LstdFlags|l.Lshortfile)
	}

	return &pubStore{
		hostname:        ctx.Hostname(),
		maxRetries:      maxRetries,
		idleTimeout:     idleTimeout,
		pubPoolsCapcity: poolCapcity,
		pubPools:        make(map[string]*pubPool),
		jobPools:        make(map[string]*jobPool),
		wg:              wg,
		dryRun:          dryRun,
		shutdownCh:      make(chan struct{}),
	}
}

func (this *pubStore) Name() string {
	return "kafka"
}

func (this *pubStore) Start() (err error) {
	this.wg.Add(1)
	defer this.wg.Done()

	// warmup: create pools according the current kafka topology
	for _, cluster := range meta.Default.ClusterNames() {
		this.pubPools[cluster] = newPubPool(this, cluster,
			meta.Default.BrokerList(cluster), this.pubPoolsCapcity)
	}

	this.refreshJobPoolNodes()

	go func() {
		for {
			select {
			case <-meta.Default.RefreshEvent():
				this.doRefresh()

			case <-this.shutdownCh:
				log.Trace("pub store[%s] stopped", this.Name())
				return
			}
		}
	}()

	return
}

func (this *pubStore) Stop() {
	this.pubPoolsLock.Lock()
	defer this.pubPoolsLock.Unlock()

	// close all kafka connections
	for _, pool := range this.pubPools {
		pool.Close()
	}

	close(this.shutdownCh)
}

func (this *pubStore) refreshJobPoolNodes() {
	if disqueAddrs, err := meta.Default.KatewayDisqueAddrs(); err == nil {
		log.Debug("disques: %+v", disqueAddrs)

		for cluster, addrs := range disqueAddrs {
			if _, present := this.jobPools[cluster]; !present {
				// found a new cluster of disque
				this.jobPools[cluster] = newJobPool(addrs)
				if e := this.jobPools[cluster].RefreshNodes(); e != nil {
					log.Error("disque[%s] refresh nodes: %v", cluster, e)

					// unload this problemetic cluster
					delete(this.jobPools, cluster)
				}
			} else {
				this.jobPools[cluster].RefreshNodes()
			}
		}
	} else {
		// just log, still using the current pools
		log.Error("disque addrs fetch: %v", err)
	}
}

func (this *pubStore) doRefresh() {
	// TODO the lock is too big, should consider cluster level refresh
	this.pubPoolsLock.Lock()
	defer this.pubPoolsLock.Unlock()
	this.jobPoolsLock.Lock()
	defer this.jobPoolsLock.Unlock()

	if time.Since(this.lastRefreshedAt) <= time.Second*5 {
		log.Warn("ignored too frequent refresh: %s", time.Since(this.lastRefreshedAt))
		return
	}

	// job pools
	this.refreshJobPoolNodes()

	// pub pool
	activeClusters := make(map[string]struct{})
	for _, cluster := range meta.Default.ClusterNames() {
		activeClusters[cluster] = struct{}{}
		if _, present := this.pubPools[cluster]; !present {
			// found a new cluster
			this.pubPools[cluster] = newPubPool(this, cluster,
				meta.Default.BrokerList(cluster), this.pubPoolsCapcity)
		} else {
			this.pubPools[cluster].RefreshBrokerList(meta.Default.BrokerList(cluster))
		}
	}

	// shutdown the dead clusters
	for cluster, pool := range this.pubPools {
		if _, present := activeClusters[cluster]; !present {
			// this cluster is dead or removed forever
			pool.Close()
			delete(this.pubPools, cluster)
		}
	}

	this.lastRefreshedAt = time.Now()
}
<|endoftext|>