<commit_before>/*
Copyright 2018 The Rook Authors. All rights reserved.

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

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

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

// Package discover to discover devices on storage nodes.
package discover

import (
	"encoding/json"
	"fmt"
	"os"
	"regexp"

	"github.com/coreos/pkg/capnslog"
	rookalpha "github.com/rook/rook/pkg/apis/rook.io/v1alpha2"
	"github.com/rook/rook/pkg/clusterd"
	discoverDaemon "github.com/rook/rook/pkg/daemon/discover"
	"github.com/rook/rook/pkg/operator/k8sutil"
	"github.com/rook/rook/pkg/util/sys"

	"k8s.io/api/core/v1"
	extensions "k8s.io/api/extensions/v1beta1"
	"k8s.io/api/rbac/v1beta1"
	kserrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
)

const (
	discoverDaemonsetName             = "rook-discover"
	discoverDaemonsetTolerationEnv    = "DISCOVER_TOLERATION"
	discoverDaemonsetTolerationKeyEnv = "DISCOVER_TOLERATION_KEY"
)

var logger = capnslog.NewPackageLogger("github.com/rook/rook", "op-discover")

var accessRules = []v1beta1.PolicyRule{
	{
		APIGroups: []string{""},
		Resources: []string{"configmaps"},
		Verbs:     []string{"get", "list", "update", "create", "delete"},
	},
}

// Discover reference to be deployed
type Discover struct {
	clientset kubernetes.Interface
}

// New creates an instance of Discover
func New(clientset kubernetes.Interface) *Discover {
	return &Discover{
		clientset: clientset,
	}
}

// Start the discover
func (d *Discover) Start(namespace, discoverImage string) error {

	err := k8sutil.MakeRole(d.clientset, namespace, discoverDaemonsetName, accessRules, nil)
	if err != nil {
		return fmt.Errorf("failed to init RBAC for rook-discover. %+v", err)
	}

	err = d.createDiscoverDaemonSet(namespace, discoverImage)
	if err != nil {
		return fmt.Errorf("Error starting discover daemonset: %v", err)
	}
	return nil
}

func (d *Discover) createDiscoverDaemonSet(namespace, discoverImage string) error {
	privileged := false
	ds := &extensions.DaemonSet{
		ObjectMeta: metav1.ObjectMeta{
			Name: discoverDaemonsetName,
		},
		Spec: extensions.DaemonSetSpec{
			UpdateStrategy: extensions.DaemonSetUpdateStrategy{
				Type: extensions.RollingUpdateDaemonSetStrategyType,
			},
			Template: v1.PodTemplateSpec{
				ObjectMeta: metav1.ObjectMeta{
					Labels: map[string]string{
						"app": discoverDaemonsetName,
					},
				},
				Spec: v1.PodSpec{
					ServiceAccountName: discoverDaemonsetName,
					Containers: []v1.Container{
						{
							Name:  discoverDaemonsetName,
							Image: discoverImage,
							Args:  []string{"discover"},
							SecurityContext: &v1.SecurityContext{
								Privileged: &privileged,
							},
							VolumeMounts: []v1.VolumeMount{
								{
									Name:      "dev",
									MountPath: "/dev",
									ReadOnly:  true,
								},
								{
									Name:      "sys",
									MountPath: "/sys",
									ReadOnly:  true,
								},
								{
									Name:      "udev",
									MountPath: "/run/udev",
									ReadOnly:  true,
								},
							},
							Env: []v1.EnvVar{
								k8sutil.NamespaceEnvVar(),
								k8sutil.NodeEnvVar(),
							},
						},
					},
					Volumes: []v1.Volume{
						{
							Name: "dev",
							VolumeSource: v1.VolumeSource{
								HostPath: &v1.HostPathVolumeSource{
									Path: "/dev",
								},
							},
						},
						{
							Name: "sys",
							VolumeSource: v1.VolumeSource{
								HostPath: &v1.HostPathVolumeSource{
									Path: "/sys",
								},
							},
						},
						{
							Name: "udev",
							VolumeSource: v1.VolumeSource{
								HostPath: &v1.HostPathVolumeSource{
									Path: "/run/udev",
								},
							},
						},
					},
					HostNetwork: false,
				},
			},
		},
	}

	// Add toleration if any
	tolerationValue := os.Getenv(discoverDaemonsetTolerationEnv)
	if tolerationValue != "" {
		ds.Spec.Template.Spec.Tolerations = []v1.Toleration{
			{
				Effect:   v1.TaintEffect(tolerationValue),
				Operator: v1.TolerationOpExists,
				Key:      os.Getenv(discoverDaemonsetTolerationKeyEnv),
			},
		}
	}

	_, err := d.clientset.Extensions().DaemonSets(namespace).Create(ds)
	if err != nil {
		if !kserrors.IsAlreadyExists(err) {
			return fmt.Errorf("failed to create rook-discover daemon set. %+v", err)
		}
		logger.Infof("rook-discover daemonset already exists")
	} else {
		logger.Infof("rook-discover daemonset started")
	}
	return nil

}

func ListDevices(context *clusterd.Context, namespace, nodeName string) (map[string][]sys.LocalDisk, error) {
	var devices map[string][]sys.LocalDisk
	listOpts := metav1.ListOptions{LabelSelector: fmt.Sprintf("%s=%s", k8sutil.AppAttr, discoverDaemon.AppName)}
	cms, err := context.Clientset.CoreV1().ConfigMaps(namespace).List(listOpts)
	if err != nil {
		return devices, fmt.Errorf("failed to list device configmaps: %+v", err)
	}
	devices = make(map[string][]sys.LocalDisk, len(cms.Items))
	for _, cm := range cms.Items {
		node := cm.ObjectMeta.Labels[discoverDaemon.NodeAttr]
		if len(nodeName) > 0 && node != nodeName {
			continue
		}
		deviceJson := cm.Data[discoverDaemon.LocalDiskCMData]
		logger.Debugf("node %s, device %s", node, deviceJson)

		if len(node) == 0 || len(deviceJson) == 0 {
			continue
		}
		var d []sys.LocalDisk
		err = json.Unmarshal([]byte(deviceJson), &d)
		if err != nil {
			logger.Warningf("failed to unmarshal %s", deviceJson)
			continue
		}
		devices[node] = d
	}
	logger.Debugf("devices %+v", devices)
	return devices, nil
}

func GetAvailableDevices(context *clusterd.Context, nodeName, clusterName string, devices []rookalpha.Device, filter string, useAllDevices bool) ([]rookalpha.Device, error) {
	results := []rookalpha.Device{}
	if len(devices) == 0 && len(filter) == 0 && !useAllDevices {
		return results, nil
	}
	namespace := os.Getenv(k8sutil.PodNamespaceEnvVar)
	allDevices, err := ListDevices(context, namespace, nodeName)
	if err != nil {
		return results, err
	}
	nodeDevices, ok := allDevices[nodeName]
	if !ok {
		return results, fmt.Errorf("node %s has no devices", nodeName)
	}
	if len(devices) > 0 {
		for i := range devices {
			for j := range nodeDevices {
				if devices[i].Name == nodeDevices[j].Name {
					results = append(results, devices[i])
				}
			}
		}
	} else if len(filter) >= 0 {
		for i := range nodeDevices {
			//TODO support filter based on other keys
			matched, err := regexp.Match(filter, []byte(nodeDevices[i].Name))
			if err == nil && matched {
				d := rookalpha.Device{
					Name: nodeDevices[i].Name,
				}
				results = append(results, d)
			}
		}
	} else if useAllDevices {
		for i := range nodeDevices {
			d := rookalpha.Device{
				Name: nodeDevices[i].Name,
			}
			results = append(results, d)
		}
	}

	return results, nil
}
<commit_msg>Update Rook Discover DaemonSet when it already exists<commit_after>/*
Copyright 2018 The Rook Authors. All rights reserved.

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

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

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

// Package discover to discover devices on storage nodes.
package discover

import (
	"encoding/json"
	"fmt"
	"os"
	"regexp"

	"github.com/coreos/pkg/capnslog"
	rookalpha "github.com/rook/rook/pkg/apis/rook.io/v1alpha2"
	"github.com/rook/rook/pkg/clusterd"
	discoverDaemon "github.com/rook/rook/pkg/daemon/discover"
	"github.com/rook/rook/pkg/operator/k8sutil"
	"github.com/rook/rook/pkg/util/sys"

	"k8s.io/api/core/v1"
	extensions "k8s.io/api/extensions/v1beta1"
	"k8s.io/api/rbac/v1beta1"
	kserrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
)

const (
	discoverDaemonsetName             = "rook-discover"
	discoverDaemonsetTolerationEnv    = "DISCOVER_TOLERATION"
	discoverDaemonsetTolerationKeyEnv = "DISCOVER_TOLERATION_KEY"
)

var logger = capnslog.NewPackageLogger("github.com/rook/rook", "op-discover")

var accessRules = []v1beta1.PolicyRule{
	{
		APIGroups: []string{""},
		Resources: []string{"configmaps"},
		Verbs:     []string{"get", "list", "update", "create", "delete"},
	},
}

// Discover reference to be deployed
type Discover struct {
	clientset kubernetes.Interface
}

// New creates an instance of Discover
func New(clientset kubernetes.Interface) *Discover {
	return &Discover{
		clientset: clientset,
	}
}

// Start the discover
func (d *Discover) Start(namespace, discoverImage string) error {

	err := k8sutil.MakeRole(d.clientset, namespace, discoverDaemonsetName, accessRules, nil)
	if err != nil {
		return fmt.Errorf("failed to init RBAC for rook-discover. %+v", err)
	}

	err = d.createDiscoverDaemonSet(namespace, discoverImage)
	if err != nil {
		return fmt.Errorf("Error starting discover daemonset: %v", err)
	}
	return nil
}

func (d *Discover) createDiscoverDaemonSet(namespace, discoverImage string) error {
	privileged := false
	ds := &extensions.DaemonSet{
		ObjectMeta: metav1.ObjectMeta{
			Name: discoverDaemonsetName,
		},
		Spec: extensions.DaemonSetSpec{
			UpdateStrategy: extensions.DaemonSetUpdateStrategy{
				Type: extensions.RollingUpdateDaemonSetStrategyType,
			},
			Template: v1.PodTemplateSpec{
				ObjectMeta: metav1.ObjectMeta{
					Labels: map[string]string{
						"app": discoverDaemonsetName,
					},
				},
				Spec: v1.PodSpec{
					ServiceAccountName: discoverDaemonsetName,
					Containers: []v1.Container{
						{
							Name:  discoverDaemonsetName,
							Image: discoverImage,
							Args:  []string{"discover"},
							SecurityContext: &v1.SecurityContext{
								Privileged: &privileged,
							},
							VolumeMounts: []v1.VolumeMount{
								{
									Name:      "dev",
									MountPath: "/dev",
									ReadOnly:  true,
								},
								{
									Name:      "sys",
									MountPath: "/sys",
									ReadOnly:  true,
								},
								{
									Name:      "udev",
									MountPath: "/run/udev",
									ReadOnly:  true,
								},
							},
							Env: []v1.EnvVar{
								k8sutil.NamespaceEnvVar(),
								k8sutil.NodeEnvVar(),
							},
						},
					},
					Volumes: []v1.Volume{
						{
							Name: "dev",
							VolumeSource: v1.VolumeSource{
								HostPath: &v1.HostPathVolumeSource{
									Path: "/dev",
								},
							},
						},
						{
							Name: "sys",
							VolumeSource: v1.VolumeSource{
								HostPath: &v1.HostPathVolumeSource{
									Path: "/sys",
								},
							},
						},
						{
							Name: "udev",
							VolumeSource: v1.VolumeSource{
								HostPath: &v1.HostPathVolumeSource{
									Path: "/run/udev",
								},
							},
						},
					},
					HostNetwork: false,
				},
			},
		},
	}

	// Add toleration if any
	tolerationValue := os.Getenv(discoverDaemonsetTolerationEnv)
	if tolerationValue != "" {
		ds.Spec.Template.Spec.Tolerations = []v1.Toleration{
			{
				Effect:   v1.TaintEffect(tolerationValue),
				Operator: v1.TolerationOpExists,
				Key:      os.Getenv(discoverDaemonsetTolerationKeyEnv),
			},
		}
	}

	_, err := d.clientset.Extensions().DaemonSets(namespace).Create(ds)
	if err != nil {
		if !kserrors.IsAlreadyExists(err) {
			return fmt.Errorf("failed to create rook-discover daemon set. %+v", err)
		}
		logger.Infof("rook-discover daemonset already exists, updating ...")
		_, err = d.clientset.Extensions().DaemonSets(namespace).Update(ds)
		if err != nil {
			return fmt.Errorf("failed to update rook-discover daemon set. %+v", err)
		}
	} else {
		logger.Infof("rook-discover daemonset started")
	}
	return nil

}

func ListDevices(context *clusterd.Context, namespace, nodeName string) (map[string][]sys.LocalDisk, error) {
	var devices map[string][]sys.LocalDisk
	listOpts := metav1.ListOptions{LabelSelector: fmt.Sprintf("%s=%s", k8sutil.AppAttr, discoverDaemon.AppName)}
	cms, err := context.Clientset.CoreV1().ConfigMaps(namespace).List(listOpts)
	if err != nil {
		return devices, fmt.Errorf("failed to list device configmaps: %+v", err)
	}
	devices = make(map[string][]sys.LocalDisk, len(cms.Items))
	for _, cm := range cms.Items {
		node := cm.ObjectMeta.Labels[discoverDaemon.NodeAttr]
		if len(nodeName) > 0 && node != nodeName {
			continue
		}
		deviceJson := cm.Data[discoverDaemon.LocalDiskCMData]
		logger.Debugf("node %s, device %s", node, deviceJson)

		if len(node) == 0 || len(deviceJson) == 0 {
			continue
		}
		var d []sys.LocalDisk
		err = json.Unmarshal([]byte(deviceJson), &d)
		if err != nil {
			logger.Warningf("failed to unmarshal %s", deviceJson)
			continue
		}
		devices[node] = d
	}
	logger.Debugf("devices %+v", devices)
	return devices, nil
}

func GetAvailableDevices(context *clusterd.Context, nodeName, clusterName string, devices []rookalpha.Device, filter string, useAllDevices bool) ([]rookalpha.Device, error) {
	results := []rookalpha.Device{}
	if len(devices) == 0 && len(filter) == 0 && !useAllDevices {
		return results, nil
	}
	namespace := os.Getenv(k8sutil.PodNamespaceEnvVar)
	allDevices, err := ListDevices(context, namespace, nodeName)
	if err != nil {
		return results, err
	}
	nodeDevices, ok := allDevices[nodeName]
	if !ok {
		return results, fmt.Errorf("node %s has no devices", nodeName)
	}
	if len(devices) > 0 {
		for i := range devices {
			for j := range nodeDevices {
				if devices[i].Name == nodeDevices[j].Name {
					results = append(results, devices[i])
				}
			}
		}
	} else if len(filter) >= 0 {
		for i := range nodeDevices {
			//TODO support filter based on other keys
			matched, err := regexp.Match(filter, []byte(nodeDevices[i].Name))
			if err == nil && matched {
				d := rookalpha.Device{
					Name: nodeDevices[i].Name,
				}
				results = append(results, d)
			}
		}
	} else if useAllDevices {
		for i := range nodeDevices {
			d := rookalpha.Device{
				Name: nodeDevices[i].Name,
			}
			results = append(results, d)
		}
	}

	return results, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Knative Authors

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

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

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

package testing

import (
	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	apiextensionsv1beta1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
	fakeapiextensionsclientset "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset/fake"
	apiextensionsv1beta1listers "k8s.io/apiextensions-apiserver/pkg/client/listers/apiextensions/v1beta1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	fakekubeclientset "k8s.io/client-go/kubernetes/fake"
	appsv1listers "k8s.io/client-go/listers/apps/v1"
	corev1listers "k8s.io/client-go/listers/core/v1"
	rbacv1listers "k8s.io/client-go/listers/rbac/v1"
	"k8s.io/client-go/tools/cache"
	configsv1alpha1 "knative.dev/eventing/pkg/apis/configs/v1alpha1"
	eventingv1alpha1 "knative.dev/eventing/pkg/apis/eventing/v1alpha1"
	eventingv1beta1 "knative.dev/eventing/pkg/apis/eventing/v1beta1"
	flowsv1alpha1 "knative.dev/eventing/pkg/apis/flows/v1alpha1"
	messagingv1alpha1 "knative.dev/eventing/pkg/apis/messaging/v1alpha1"
	sourcesv1alpha1 "knative.dev/eventing/pkg/apis/sources/v1alpha1"
	sourcesv1alpha2 "knative.dev/eventing/pkg/apis/sources/v1alpha2"
	fakeeventingclientset "knative.dev/eventing/pkg/client/clientset/versioned/fake"
	configslisters "knative.dev/eventing/pkg/client/listers/configs/v1alpha1"
	eventinglisters "knative.dev/eventing/pkg/client/listers/eventing/v1alpha1"
	eventingv1beta1listers "knative.dev/eventing/pkg/client/listers/eventing/v1beta1"
	flowslisters "knative.dev/eventing/pkg/client/listers/flows/v1alpha1"
	messaginglisters "knative.dev/eventing/pkg/client/listers/messaging/v1alpha1"
	sourcelisters "knative.dev/eventing/pkg/client/listers/sources/v1alpha1"
	sourcev1alpha2listers "knative.dev/eventing/pkg/client/listers/sources/v1alpha2"
	duckv1 "knative.dev/pkg/apis/duck/v1"
	"knative.dev/pkg/reconciler/testing"
)

var subscriberAddToScheme = func(scheme *runtime.Scheme) error {
	scheme.AddKnownTypeWithName(schema.GroupVersionKind{Group: "testing.eventing.knative.dev", Version: "v1alpha1", Kind: "Subscriber"}, &unstructured.Unstructured{})
	return nil
}

var sourceAddToScheme = func(scheme *runtime.Scheme) error {
	scheme.AddKnownTypeWithName(schema.GroupVersionKind{Group: "testing.sources.knative.dev", Version: "v1alpha1", Kind: "TestSource"}, &duckv1.Source{})
	return nil
}

var clientSetSchemes = []func(*runtime.Scheme) error{
	fakekubeclientset.AddToScheme,
	fakeeventingclientset.AddToScheme,
	fakeapiextensionsclientset.AddToScheme,
	subscriberAddToScheme,
	sourceAddToScheme,
}

type Listers struct {
	sorter testing.ObjectSorter
}

func NewScheme() *runtime.Scheme {
	scheme := runtime.NewScheme()

	for _, addTo := range clientSetSchemes {
		addTo(scheme)
	}
	return scheme
}

func NewListers(objs []runtime.Object) Listers {
	scheme := runtime.NewScheme()

	for _, addTo := range clientSetSchemes {
		addTo(scheme)
	}

	ls := Listers{
		sorter: testing.NewObjectSorter(scheme),
	}

	ls.sorter.AddObjects(objs...)

	return ls
}

func (l *Listers) indexerFor(obj runtime.Object) cache.Indexer {
	return l.sorter.IndexerForObjectType(obj)
}

func (l *Listers) GetKubeObjects() []runtime.Object {
	return l.sorter.ObjectsForSchemeFunc(fakekubeclientset.AddToScheme)
}

func (l *Listers) GetEventingObjects() []runtime.Object {
	return l.sorter.ObjectsForSchemeFunc(fakeeventingclientset.AddToScheme)
}

func (l *Listers) GetSubscriberObjects() []runtime.Object {
	return l.sorter.ObjectsForSchemeFunc(subscriberAddToScheme)
}

func (l *Listers) GetAllObjects() []runtime.Object {
	all := l.GetSubscriberObjects()
	all = append(all, l.GetEventingObjects()...)
	all = append(all, l.GetKubeObjects()...)
	return all
}

func (l *Listers) GetSubscriptionLister() messaginglisters.SubscriptionLister {
	return messaginglisters.NewSubscriptionLister(l.indexerFor(&messagingv1alpha1.Subscription{}))
}

func (l *Listers) GetFlowsSequenceLister() flowslisters.SequenceLister {
	return flowslisters.NewSequenceLister(l.indexerFor(&flowsv1alpha1.Sequence{}))
}

func (l *Listers) GetTriggerLister() eventinglisters.TriggerLister {
	return eventinglisters.NewTriggerLister(l.indexerFor(&eventingv1alpha1.Trigger{}))
}

func (l *Listers) GetBrokerLister() eventinglisters.BrokerLister {
	return eventinglisters.NewBrokerLister(l.indexerFor(&eventingv1alpha1.Broker{}))
}

func (l *Listers) GetV1Beta1BrokerLister() eventingv1beta1listers.BrokerLister {
	return eventingv1beta1listers.NewBrokerLister(l.indexerFor(&eventingv1beta1.Broker{}))
}

func (l *Listers) GetEventTypeLister() eventinglisters.EventTypeLister {
	return eventinglisters.NewEventTypeLister(l.indexerFor(&eventingv1alpha1.EventType{}))
}

func (l *Listers) GetInMemoryChannelLister() messaginglisters.InMemoryChannelLister {
	return messaginglisters.NewInMemoryChannelLister(l.indexerFor(&messagingv1alpha1.InMemoryChannel{}))
}

func (l *Listers) GetMessagingChannelLister() messaginglisters.ChannelLister {
	return messaginglisters.NewChannelLister(l.indexerFor(&messagingv1alpha1.Channel{}))
}

func (l *Listers) GetFlowsParallelLister() flowslisters.ParallelLister {
	return flowslisters.NewParallelLister(l.indexerFor(&flowsv1alpha1.Parallel{}))
}

func (l *Listers) GetApiServerSourceLister() sourcelisters.ApiServerSourceLister {
	return sourcelisters.NewApiServerSourceLister(l.indexerFor(&sourcesv1alpha1.ApiServerSource{}))
}

func (l *Listers) GetPingSourceLister() sourcelisters.PingSourceLister {
	return sourcelisters.NewPingSourceLister(l.indexerFor(&sourcesv1alpha1.PingSource{}))
}

func (l *Listers) GetSinkBindingLister() sourcelisters.SinkBindingLister {
	return sourcelisters.NewSinkBindingLister(l.indexerFor(&sourcesv1alpha1.SinkBinding{}))
}

func (l *Listers) GetPingSourceV1alpha2Lister() sourcev1alpha2listers.PingSourceLister {
	return sourcev1alpha2listers.NewPingSourceLister(l.indexerFor(&sourcesv1alpha2.PingSource{}))
}

func (l *Listers) GetContainerSourceLister() sourcev1alpha2listers.ContainerSourceLister {
	return sourcev1alpha2listers.NewContainerSourceLister(l.indexerFor(&sourcesv1alpha2.ContainerSource{}))
}

func (l *Listers) GetSinkBindingV1alpha2Lister() sourcev1alpha2listers.SinkBindingLister {
	return sourcev1alpha2listers.NewSinkBindingLister(l.indexerFor(&sourcesv1alpha2.SinkBinding{}))
}

func (l *Listers) GetApiServerSourceV1alpha2Lister() sourcev1alpha2listers.ApiServerSourceLister {
	return sourcev1alpha2listers.NewApiServerSourceLister(l.indexerFor(&sourcesv1alpha2.ApiServerSource{}))
}

func (l *Listers) GetDeploymentLister() appsv1listers.DeploymentLister {
	return appsv1listers.NewDeploymentLister(l.indexerFor(&appsv1.Deployment{}))
}

func (l *Listers) GetK8sServiceLister() corev1listers.ServiceLister {
	return corev1listers.NewServiceLister(l.indexerFor(&corev1.Service{}))
}

func (l *Listers) GetNamespaceLister() corev1listers.NamespaceLister {
	return corev1listers.NewNamespaceLister(l.indexerFor(&corev1.Namespace{}))
}

func (l *Listers) GetServiceAccountLister() corev1listers.ServiceAccountLister {
	return corev1listers.NewServiceAccountLister(l.indexerFor(&corev1.ServiceAccount{}))
}

func (l *Listers) GetServiceLister() corev1listers.ServiceLister {
	return corev1listers.NewServiceLister(l.indexerFor(&corev1.Service{}))
}

func (l *Listers) GetRoleBindingLister() rbacv1listers.RoleBindingLister {
	return rbacv1listers.NewRoleBindingLister(l.indexerFor(&rbacv1.RoleBinding{}))
}

func (l *Listers) GetEndpointsLister() corev1listers.EndpointsLister {
	return corev1listers.NewEndpointsLister(l.indexerFor(&corev1.Endpoints{}))
}

func (l *Listers) GetConfigMapLister() corev1listers.ConfigMapLister {
	return corev1listers.NewConfigMapLister(l.indexerFor(&corev1.ConfigMap{}))
}

func (l *Listers) GetCustomResourceDefinitionLister() apiextensionsv1beta1listers.CustomResourceDefinitionLister {
	return apiextensionsv1beta1listers.NewCustomResourceDefinitionLister(l.indexerFor(&apiextensionsv1beta1.CustomResourceDefinition{}))
}

func (l *Listers) GetConfigMapPropagationLister() configslisters.ConfigMapPropagationLister {
	return configslisters.NewConfigMapPropagationLister(l.indexerFor(&configsv1alpha1.ConfigMapPropagation{}))
}
<commit_msg>Add PodLister to reconciler test listers (#3030)<commit_after>/*
Copyright 2018 The Knative Authors

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

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

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

package testing

import (
	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	apiextensionsv1beta1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
	fakeapiextensionsclientset "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset/fake"
	apiextensionsv1beta1listers "k8s.io/apiextensions-apiserver/pkg/client/listers/apiextensions/v1beta1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	fakekubeclientset "k8s.io/client-go/kubernetes/fake"
	appsv1listers "k8s.io/client-go/listers/apps/v1"
	corev1listers "k8s.io/client-go/listers/core/v1"
	rbacv1listers "k8s.io/client-go/listers/rbac/v1"
	"k8s.io/client-go/tools/cache"
	configsv1alpha1 "knative.dev/eventing/pkg/apis/configs/v1alpha1"
	eventingv1alpha1 "knative.dev/eventing/pkg/apis/eventing/v1alpha1"
	eventingv1beta1 "knative.dev/eventing/pkg/apis/eventing/v1beta1"
	flowsv1alpha1 "knative.dev/eventing/pkg/apis/flows/v1alpha1"
	messagingv1alpha1 "knative.dev/eventing/pkg/apis/messaging/v1alpha1"
	sourcesv1alpha1 "knative.dev/eventing/pkg/apis/sources/v1alpha1"
	sourcesv1alpha2 "knative.dev/eventing/pkg/apis/sources/v1alpha2"
	fakeeventingclientset "knative.dev/eventing/pkg/client/clientset/versioned/fake"
	configslisters "knative.dev/eventing/pkg/client/listers/configs/v1alpha1"
	eventinglisters "knative.dev/eventing/pkg/client/listers/eventing/v1alpha1"
	eventingv1beta1listers "knative.dev/eventing/pkg/client/listers/eventing/v1beta1"
	flowslisters "knative.dev/eventing/pkg/client/listers/flows/v1alpha1"
	messaginglisters "knative.dev/eventing/pkg/client/listers/messaging/v1alpha1"
	sourcelisters "knative.dev/eventing/pkg/client/listers/sources/v1alpha1"
	sourcev1alpha2listers "knative.dev/eventing/pkg/client/listers/sources/v1alpha2"
	duckv1 "knative.dev/pkg/apis/duck/v1"
	"knative.dev/pkg/reconciler/testing"
)

var subscriberAddToScheme = func(scheme *runtime.Scheme) error {
	scheme.AddKnownTypeWithName(schema.GroupVersionKind{Group: "testing.eventing.knative.dev", Version: "v1alpha1", Kind: "Subscriber"}, &unstructured.Unstructured{})
	return nil
}

var sourceAddToScheme = func(scheme *runtime.Scheme) error {
	scheme.AddKnownTypeWithName(schema.GroupVersionKind{Group: "testing.sources.knative.dev", Version: "v1alpha1", Kind: "TestSource"}, &duckv1.Source{})
	return nil
}

var clientSetSchemes = []func(*runtime.Scheme) error{
	fakekubeclientset.AddToScheme,
	fakeeventingclientset.AddToScheme,
	fakeapiextensionsclientset.AddToScheme,
	subscriberAddToScheme,
	sourceAddToScheme,
}

type Listers struct {
	sorter testing.ObjectSorter
}

func NewScheme() *runtime.Scheme {
	scheme := runtime.NewScheme()

	for _, addTo := range clientSetSchemes {
		addTo(scheme)
	}
	return scheme
}

func NewListers(objs []runtime.Object) Listers {
	scheme := runtime.NewScheme()

	for _, addTo := range clientSetSchemes {
		addTo(scheme)
	}

	ls := Listers{
		sorter: testing.NewObjectSorter(scheme),
	}

	ls.sorter.AddObjects(objs...)

	return ls
}

func (l *Listers) indexerFor(obj runtime.Object) cache.Indexer {
	return l.sorter.IndexerForObjectType(obj)
}

func (l *Listers) GetKubeObjects() []runtime.Object {
	return l.sorter.ObjectsForSchemeFunc(fakekubeclientset.AddToScheme)
}

func (l *Listers) GetEventingObjects() []runtime.Object {
	return l.sorter.ObjectsForSchemeFunc(fakeeventingclientset.AddToScheme)
}

func (l *Listers) GetSubscriberObjects() []runtime.Object {
	return l.sorter.ObjectsForSchemeFunc(subscriberAddToScheme)
}

func (l *Listers) GetAllObjects() []runtime.Object {
	all := l.GetSubscriberObjects()
	all = append(all, l.GetEventingObjects()...)
	all = append(all, l.GetKubeObjects()...)
	return all
}

func (l *Listers) GetSubscriptionLister() messaginglisters.SubscriptionLister {
	return messaginglisters.NewSubscriptionLister(l.indexerFor(&messagingv1alpha1.Subscription{}))
}

func (l *Listers) GetFlowsSequenceLister() flowslisters.SequenceLister {
	return flowslisters.NewSequenceLister(l.indexerFor(&flowsv1alpha1.Sequence{}))
}

func (l *Listers) GetTriggerLister() eventinglisters.TriggerLister {
	return eventinglisters.NewTriggerLister(l.indexerFor(&eventingv1alpha1.Trigger{}))
}

func (l *Listers) GetBrokerLister() eventinglisters.BrokerLister {
	return eventinglisters.NewBrokerLister(l.indexerFor(&eventingv1alpha1.Broker{}))
}

func (l *Listers) GetV1Beta1BrokerLister() eventingv1beta1listers.BrokerLister {
	return eventingv1beta1listers.NewBrokerLister(l.indexerFor(&eventingv1beta1.Broker{}))
}

func (l *Listers) GetEventTypeLister() eventinglisters.EventTypeLister {
	return eventinglisters.NewEventTypeLister(l.indexerFor(&eventingv1alpha1.EventType{}))
}

func (l *Listers) GetInMemoryChannelLister() messaginglisters.InMemoryChannelLister {
	return messaginglisters.NewInMemoryChannelLister(l.indexerFor(&messagingv1alpha1.InMemoryChannel{}))
}

func (l *Listers) GetMessagingChannelLister() messaginglisters.ChannelLister {
	return messaginglisters.NewChannelLister(l.indexerFor(&messagingv1alpha1.Channel{}))
}

func (l *Listers) GetFlowsParallelLister() flowslisters.ParallelLister {
	return flowslisters.NewParallelLister(l.indexerFor(&flowsv1alpha1.Parallel{}))
}

func (l *Listers) GetApiServerSourceLister() sourcelisters.ApiServerSourceLister {
	return sourcelisters.NewApiServerSourceLister(l.indexerFor(&sourcesv1alpha1.ApiServerSource{}))
}

func (l *Listers) GetPingSourceLister() sourcelisters.PingSourceLister {
	return sourcelisters.NewPingSourceLister(l.indexerFor(&sourcesv1alpha1.PingSource{}))
}

func (l *Listers) GetSinkBindingLister() sourcelisters.SinkBindingLister {
	return sourcelisters.NewSinkBindingLister(l.indexerFor(&sourcesv1alpha1.SinkBinding{}))
}

func (l *Listers) GetPingSourceV1alpha2Lister() sourcev1alpha2listers.PingSourceLister {
	return sourcev1alpha2listers.NewPingSourceLister(l.indexerFor(&sourcesv1alpha2.PingSource{}))
}

func (l *Listers) GetContainerSourceLister() sourcev1alpha2listers.ContainerSourceLister {
	return sourcev1alpha2listers.NewContainerSourceLister(l.indexerFor(&sourcesv1alpha2.ContainerSource{}))
}

func (l *Listers) GetSinkBindingV1alpha2Lister() sourcev1alpha2listers.SinkBindingLister {
	return sourcev1alpha2listers.NewSinkBindingLister(l.indexerFor(&sourcesv1alpha2.SinkBinding{}))
}

func (l *Listers) GetApiServerSourceV1alpha2Lister() sourcev1alpha2listers.ApiServerSourceLister {
	return sourcev1alpha2listers.NewApiServerSourceLister(l.indexerFor(&sourcesv1alpha2.ApiServerSource{}))
}

func (l *Listers) GetDeploymentLister() appsv1listers.DeploymentLister {
	return appsv1listers.NewDeploymentLister(l.indexerFor(&appsv1.Deployment{}))
}

func (l *Listers) GetK8sServiceLister() corev1listers.ServiceLister {
	return corev1listers.NewServiceLister(l.indexerFor(&corev1.Service{}))
}

func (l *Listers) GetNamespaceLister() corev1listers.NamespaceLister {
	return corev1listers.NewNamespaceLister(l.indexerFor(&corev1.Namespace{}))
}

func (l *Listers) GetServiceAccountLister() corev1listers.ServiceAccountLister {
	return corev1listers.NewServiceAccountLister(l.indexerFor(&corev1.ServiceAccount{}))
}

func (l *Listers) GetServiceLister() corev1listers.ServiceLister {
	return corev1listers.NewServiceLister(l.indexerFor(&corev1.Service{}))
}

func (l *Listers) GetPodLister() corev1listers.PodLister {
	return corev1listers.NewPodLister(l.indexerFor(&corev1.Pod{}))
}

func (l *Listers) GetRoleBindingLister() rbacv1listers.RoleBindingLister {
	return rbacv1listers.NewRoleBindingLister(l.indexerFor(&rbacv1.RoleBinding{}))
}

func (l *Listers) GetEndpointsLister() corev1listers.EndpointsLister {
	return corev1listers.NewEndpointsLister(l.indexerFor(&corev1.Endpoints{}))
}

func (l *Listers) GetConfigMapLister() corev1listers.ConfigMapLister {
	return corev1listers.NewConfigMapLister(l.indexerFor(&corev1.ConfigMap{}))
}

func (l *Listers) GetCustomResourceDefinitionLister() apiextensionsv1beta1listers.CustomResourceDefinitionLister {
	return apiextensionsv1beta1listers.NewCustomResourceDefinitionLister(l.indexerFor(&apiextensionsv1beta1.CustomResourceDefinition{}))
}

func (l *Listers) GetConfigMapPropagationLister() configslisters.ConfigMapPropagationLister {
	return configslisters.NewConfigMapPropagationLister(l.indexerFor(&configsv1alpha1.ConfigMapPropagation{}))
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Knative Authors

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

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

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

package testing

import (
	"context"
	//	"time"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	//	"k8s.io/apimachinery/pkg/types"

	"knative.dev/pkg/apis"
	apisv1alpha1 "knative.dev/pkg/apis/v1alpha1"

	"github.com/google/knative-gcp/pkg/apis/events/v1alpha1"
)

// StorageOption enables further configuration of a Storage.
type StorageOption func(*v1alpha1.Storage)

// NewStorage creates a Storage with StorageOptions
func NewStorage(name, namespace string, so ...StorageOption) *v1alpha1.Storage {
	s := &v1alpha1.Storage{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
			UID:       "test-storage-uid",
		},
	}
	for _, opt := range so {
		opt(s)
	}
	s.SetDefaults(context.Background())
	return s
}

func WithStorageBucket(bucket string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Spec.Bucket = bucket
	}
}

func WithStorageSink(gvk metav1.GroupVersionKind, name string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Spec.Sink = apisv1alpha1.Destination{
			ObjectReference: &corev1.ObjectReference{
				APIVersion: apiVersion(gvk),
				Kind:       gvk.Kind,
				Name:       name,
			},
		}
	}
}

// WithInitStorageConditions initializes the Storages's conditions.
func WithInitStorageConditions(s *v1alpha1.Storage) {
	s.Status.InitializeConditions()
}

// WithStorageTopicNotReady marks the condition that the
// topic is not ready
func WithStorageTopicNotReady(reason, message string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.MarkTopicNotReady(reason, message)
	}
}

// WithStorageTopicNotReady marks the condition that the
// topic is not ready
func WithStorageTopicReady(topicID string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.MarkTopicReady()
		s.Status.TopicID = topicID
	}
}

// WithStoragePullSubscriptionNotReady marks the condition that the
// topic is not ready
func WithStoragePullSubscriptionNotReady(reason, message string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.MarkPullSubscriptionNotReady(reason, message)
	}
}

// WithStoragePullSubscriptionNotReady marks the condition that the
// topic is not ready
func WithStoragePullSubscriptionReady() StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.MarkPullSubscriptionReady()
	}
}

// WithStorageGCSNotReady marks the condition that the
// topic is not ready
func WithStorageGCSNotReady(reason, message string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.MarkGCSNotReady(reason, message)
	}
}

// WithStorageGCSNotReady marks the condition that the
// topic is not ready
func WithStorageGCSReady() StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.MarkGCSReady()
	}
}

// WithStorageSinkURI sets the status for sink URI
func WithStorageSinkURI(url *apis.URL) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.SinkURI = url
	}
}

// WithStorageNotificationId sets the status for Notification ID
func WithStorageNotificationID(notificationID string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.NotificationID = notificationID
	}
}

// WithStorageProjectId sets the status for Project ID
func WithStorageProjectID(notificationID string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.ProjectID = notificationID
	}
}

func WithStorageFinalizers(finalizers ...string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Finalizers = finalizers
	}
}
<commit_msg>remove unused imports<commit_after>/*
Copyright 2019 The Knative Authors

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

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

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

package testing

import (
	"context"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"knative.dev/pkg/apis"
	apisv1alpha1 "knative.dev/pkg/apis/v1alpha1"

	"github.com/google/knative-gcp/pkg/apis/events/v1alpha1"
)

// StorageOption enables further configuration of a Storage.
type StorageOption func(*v1alpha1.Storage)

// NewStorage creates a Storage with StorageOptions
func NewStorage(name, namespace string, so ...StorageOption) *v1alpha1.Storage {
	s := &v1alpha1.Storage{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
			UID:       "test-storage-uid",
		},
	}
	for _, opt := range so {
		opt(s)
	}
	s.SetDefaults(context.Background())
	return s
}

func WithStorageBucket(bucket string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Spec.Bucket = bucket
	}
}

func WithStorageSink(gvk metav1.GroupVersionKind, name string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Spec.Sink = apisv1alpha1.Destination{
			ObjectReference: &corev1.ObjectReference{
				APIVersion: apiVersion(gvk),
				Kind:       gvk.Kind,
				Name:       name,
			},
		}
	}
}

// WithInitStorageConditions initializes the Storages's conditions.
func WithInitStorageConditions(s *v1alpha1.Storage) {
	s.Status.InitializeConditions()
}

// WithStorageTopicNotReady marks the condition that the
// topic is not ready
func WithStorageTopicNotReady(reason, message string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.MarkTopicNotReady(reason, message)
	}
}

// WithStorageTopicNotReady marks the condition that the
// topic is not ready
func WithStorageTopicReady(topicID string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.MarkTopicReady()
		s.Status.TopicID = topicID
	}
}

// WithStoragePullSubscriptionNotReady marks the condition that the
// topic is not ready
func WithStoragePullSubscriptionNotReady(reason, message string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.MarkPullSubscriptionNotReady(reason, message)
	}
}

// WithStoragePullSubscriptionNotReady marks the condition that the
// topic is not ready
func WithStoragePullSubscriptionReady() StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.MarkPullSubscriptionReady()
	}
}

// WithStorageGCSNotReady marks the condition that the
// topic is not ready
func WithStorageGCSNotReady(reason, message string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.MarkGCSNotReady(reason, message)
	}
}

// WithStorageGCSNotReady marks the condition that the
// topic is not ready
func WithStorageGCSReady() StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.MarkGCSReady()
	}
}

// WithStorageSinkURI sets the status for sink URI
func WithStorageSinkURI(url *apis.URL) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.SinkURI = url
	}
}

// WithStorageNotificationId sets the status for Notification ID
func WithStorageNotificationID(notificationID string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.NotificationID = notificationID
	}
}

// WithStorageProjectId sets the status for Project ID
func WithStorageProjectID(notificationID string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Status.ProjectID = notificationID
	}
}

func WithStorageFinalizers(finalizers ...string) StorageOption {
	return func(s *v1alpha1.Storage) {
		s.Finalizers = finalizers
	}
}
<|endoftext|>


<commit_before>package podtask

import (
	"container/ring"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	log "github.com/golang/glog"
	mesos "github.com/mesos/mesos-go/mesosproto"
	"github.com/mesosphere/kubernetes-mesos/pkg/scheduler/metrics"
)

/**
HACK(jdef): we're not using etcd but k8s has implemented namespace support and
we're going to try to honor that by namespacing pod keys. Hence, the following
funcs that were stolen from:
    https://github.com/GoogleCloudPlatform/kubernetes/blob/release-0.5/pkg/registry/etcd/etcd.go
**/

const (
	PodPath                  = "/pods"
	defaultFinishedTasksSize = 1024
)

type Registry interface {
	Register(*T, error) (*T, error)
	Unregister(*T)
	Get(taskId string) (task *T, currentState StateType)
	TaskForPod(podID string) (taskID string, ok bool)
	UpdateStatus(status *mesos.TaskStatus) (*T, StateType)
	// return a list of task ID's that match the given filter, or all task ID's if filter == nil
	List(func(*T) bool) []string
}

type inMemoryRegistry struct {
	rw            sync.RWMutex
	taskRegistry  map[string]*T
	tasksFinished *ring.Ring
	podToTask     map[string]string
}

func NewInMemoryRegistry() Registry {
	return &inMemoryRegistry{
		taskRegistry:  make(map[string]*T),
		tasksFinished: ring.New(defaultFinishedTasksSize),
		podToTask:     make(map[string]string),
	}
}

func (k *inMemoryRegistry) List(accepts func(t *T) bool) (taskids []string) {
	k.rw.RLock()
	defer k.rw.RUnlock()
	for id, task := range k.taskRegistry {
		if accepts == nil || accepts(task) {
			taskids = append(taskids, id)
		}
	}
	return
}

func (k *inMemoryRegistry) TaskForPod(podID string) (taskID string, ok bool) {
	k.rw.RLock()
	defer k.rw.RUnlock()
	// assume caller is holding scheduler lock
	taskID, ok = k.podToTask[podID]
	return
}

// registers a pod task unless the spec'd error is not nil
func (k *inMemoryRegistry) Register(task *T, err error) (*T, error) {
	if err == nil {
		k.rw.Lock()
		defer k.rw.Unlock()
		k.podToTask[task.podKey] = task.ID
		k.taskRegistry[task.ID] = task
	}
	return task, err
}

func (k *inMemoryRegistry) Unregister(task *T) {
	k.rw.Lock()
	defer k.rw.Unlock()
	delete(k.podToTask, task.podKey)
	delete(k.taskRegistry, task.ID)
}

func (k *inMemoryRegistry) Get(taskId string) (*T, StateType) {
	k.rw.RLock()
	defer k.rw.RUnlock()
	return k._get(taskId)
}

// assume that the caller has already locked around access to task state
func (k *inMemoryRegistry) _get(taskId string) (*T, StateType) {
	if task, found := k.taskRegistry[taskId]; found {
		return task, task.State
	}
	return nil, StateUnknown
}

func (k *inMemoryRegistry) UpdateStatus(status *mesos.TaskStatus) (*T, StateType) {
	taskId := status.GetTaskId().GetValue()

	k.rw.Lock()
	defer k.rw.Unlock()
	task, state := k._get(taskId)

	switch status.GetState() {
	case mesos.TaskState_TASK_STAGING:
		k.handleTaskStaging(task, state, status)
	case mesos.TaskState_TASK_STARTING:
		k.handleTaskStarting(task, state, status)
	case mesos.TaskState_TASK_RUNNING:
		k.handleTaskRunning(task, state, status)
	case mesos.TaskState_TASK_FINISHED:
		k.handleTaskFinished(task, state, status)
	case mesos.TaskState_TASK_FAILED:
		k.handleTaskFailed(task, state, status)
	case mesos.TaskState_TASK_KILLED:
		k.handleTaskKilled(task, state, status)
	case mesos.TaskState_TASK_LOST:
		k.handleTaskLost(task, state, status)
	default:
		log.Warning("unhandled task status update: %+v", status)
	}
	return task, state
}

func (k *inMemoryRegistry) handleTaskStaging(task *T, state StateType, status *mesos.TaskStatus) {
	if status.GetSource() != mesos.TaskStatus_SOURCE_MASTER {
		log.Errorf("received STAGING for task %v with unexpected source: %v",
			status.GetTaskId().GetValue(), status.GetSource())
	}
}

func (k *inMemoryRegistry) handleTaskStarting(task *T, state StateType, status *mesos.TaskStatus) {
	// we expect to receive this when a launched task is finally "bound"
	// via the API server. however, there's nothing specific for us to do
	// here.
	switch state {
	case StatePending:
		task.UpdatedTime = time.Now()
		if !task.Has(Bound) {
			task.Set(Bound)
			task.bindTime = task.UpdatedTime
			timeToBind := task.bindTime.Sub(task.launchTime)
			metrics.BindLatency.Observe(metrics.InMicroseconds(timeToBind))
		}
	default:
		log.Warningf("Ignore status TASK_STARTING because the the task is not pending")
	}
}

func (k *inMemoryRegistry) handleTaskRunning(task *T, state StateType, status *mesos.TaskStatus) {
	switch state {
	case StatePending:
		task.UpdatedTime = time.Now()
		log.Infof("Received running status for pending task: %+v", status)
		fillRunningPodInfo(task, status)
		task.State = StateRunning
	case StateRunning:
		task.UpdatedTime = time.Now()
		log.V(2).Info("Ignore status TASK_RUNNING because the the task is already running")
	case StateFinished:
		log.Warningf("Ignore status TASK_RUNNING because the the task is already finished")
	default:
		log.Warningf("Ignore status TASK_RUNNING (%+v) because the the task is discarded", status.GetTaskId())
	}
}

func ParsePodStatusResult(taskStatus *mesos.TaskStatus) (result api.PodStatusResult, err error) {
	if taskStatus.Data != nil {
		err = json.Unmarshal(taskStatus.Data, &result)
	} else {
		err = fmt.Errorf("missing TaskStatus.Data")
	}
	return
}

func fillRunningPodInfo(task *T, taskStatus *mesos.TaskStatus) {
	if result, err := ParsePodStatusResult(taskStatus); err != nil {
		log.Errorf("invalid TaskStatus.Data for task '%v': %v", task.ID, err)
	} else {
		task.Pod.Status = result.Status
		log.Infof("received pod status for task %v: %+v", task.ID, task.Pod.Status)
	}
}

func (k *inMemoryRegistry) handleTaskFinished(task *T, state StateType, status *mesos.TaskStatus) {
	switch state {
	case StatePending:
		panic("Pending task finished, this couldn't happen")
	case StateRunning:
		log.V(2).Infof("received finished status for running task: %+v", status)
		delete(k.podToTask, task.podKey)
		task.State = StateFinished
		task.UpdatedTime = time.Now()
		k.tasksFinished = k.recordFinishedTask(task.ID)
	case StateFinished:
		log.Warningf("Ignore status TASK_FINISHED because the the task is already finished")
	default:
		log.Warningf("Ignore status TASK_FINISHED because the the task is not running")
	}
}

// record that a task has finished.
// older record are expunged one at a time once the historical ring buffer is saturated.
// assumes caller is holding state lock.
func (k *inMemoryRegistry) recordFinishedTask(taskId string) *ring.Ring {
	slot := k.tasksFinished.Next()
	if slot.Value != nil {
		// garbage collect older finished task from the registry
		gctaskId := slot.Value.(string)
		if gctask, found := k.taskRegistry[gctaskId]; found && gctask.State == StateFinished {
			delete(k.taskRegistry, gctaskId)
		}
	}
	slot.Value = taskId
	return slot
}

func (k *inMemoryRegistry) handleTaskFailed(task *T, state StateType, status *mesos.TaskStatus) {
	log.Errorf("task failed: %+v", status)
	switch state {
	case StatePending:
		delete(k.taskRegistry, task.ID)
		delete(k.podToTask, task.podKey)
	case StateRunning:
		delete(k.taskRegistry, task.ID)
		delete(k.podToTask, task.podKey)
	}
}

func (k *inMemoryRegistry) handleTaskKilled(task *T, state StateType, status *mesos.TaskStatus) {
	defer func() {
		msg := fmt.Sprintf("task killed: %+v, task %+v", status, task)
		if task != nil && task.Has(Deleted) {
			// we were expecting this, nothing out of the ordinary
			log.V(2).Infoln(msg)
		} else {
			log.Errorln(msg)
		}
	}()
	switch state {
	case StatePending, StateRunning:
		delete(k.taskRegistry, task.ID)
		delete(k.podToTask, task.podKey)
	}
}

func (k *inMemoryRegistry) handleTaskLost(task *T, state StateType, status *mesos.TaskStatus) {
	log.Warningf("task lost: %+v", status)
	switch state {
	case StateRunning, StatePending:
		delete(k.taskRegistry, task.ID)
		delete(k.podToTask, task.podKey)
	}
}
<commit_msg>task registry should avoid parsing taskStatus.data for reconciliation updates<commit_after>package podtask

import (
	"container/ring"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	log "github.com/golang/glog"
	mesos "github.com/mesos/mesos-go/mesosproto"
	"github.com/mesosphere/kubernetes-mesos/pkg/scheduler/metrics"
)

/**
HACK(jdef): we're not using etcd but k8s has implemented namespace support and
we're going to try to honor that by namespacing pod keys. Hence, the following
funcs that were stolen from:
    https://github.com/GoogleCloudPlatform/kubernetes/blob/release-0.5/pkg/registry/etcd/etcd.go
**/

const (
	PodPath                  = "/pods"
	defaultFinishedTasksSize = 1024
)

type Registry interface {
	Register(*T, error) (*T, error)
	Unregister(*T)
	Get(taskId string) (task *T, currentState StateType)
	TaskForPod(podID string) (taskID string, ok bool)
	UpdateStatus(status *mesos.TaskStatus) (*T, StateType)
	// return a list of task ID's that match the given filter, or all task ID's if filter == nil
	List(func(*T) bool) []string
}

type inMemoryRegistry struct {
	rw            sync.RWMutex
	taskRegistry  map[string]*T
	tasksFinished *ring.Ring
	podToTask     map[string]string
}

func NewInMemoryRegistry() Registry {
	return &inMemoryRegistry{
		taskRegistry:  make(map[string]*T),
		tasksFinished: ring.New(defaultFinishedTasksSize),
		podToTask:     make(map[string]string),
	}
}

func (k *inMemoryRegistry) List(accepts func(t *T) bool) (taskids []string) {
	k.rw.RLock()
	defer k.rw.RUnlock()
	for id, task := range k.taskRegistry {
		if accepts == nil || accepts(task) {
			taskids = append(taskids, id)
		}
	}
	return
}

func (k *inMemoryRegistry) TaskForPod(podID string) (taskID string, ok bool) {
	k.rw.RLock()
	defer k.rw.RUnlock()
	// assume caller is holding scheduler lock
	taskID, ok = k.podToTask[podID]
	return
}

// registers a pod task unless the spec'd error is not nil
func (k *inMemoryRegistry) Register(task *T, err error) (*T, error) {
	if err == nil {
		k.rw.Lock()
		defer k.rw.Unlock()
		k.podToTask[task.podKey] = task.ID
		k.taskRegistry[task.ID] = task
	}
	return task, err
}

func (k *inMemoryRegistry) Unregister(task *T) {
	k.rw.Lock()
	defer k.rw.Unlock()
	delete(k.podToTask, task.podKey)
	delete(k.taskRegistry, task.ID)
}

func (k *inMemoryRegistry) Get(taskId string) (*T, StateType) {
	k.rw.RLock()
	defer k.rw.RUnlock()
	return k._get(taskId)
}

// assume that the caller has already locked around access to task state
func (k *inMemoryRegistry) _get(taskId string) (*T, StateType) {
	if task, found := k.taskRegistry[taskId]; found {
		return task, task.State
	}
	return nil, StateUnknown
}

func (k *inMemoryRegistry) UpdateStatus(status *mesos.TaskStatus) (*T, StateType) {
	taskId := status.GetTaskId().GetValue()

	k.rw.Lock()
	defer k.rw.Unlock()
	task, state := k._get(taskId)

	switch status.GetState() {
	case mesos.TaskState_TASK_STAGING:
		k.handleTaskStaging(task, state, status)
	case mesos.TaskState_TASK_STARTING:
		k.handleTaskStarting(task, state, status)
	case mesos.TaskState_TASK_RUNNING:
		k.handleTaskRunning(task, state, status)
	case mesos.TaskState_TASK_FINISHED:
		k.handleTaskFinished(task, state, status)
	case mesos.TaskState_TASK_FAILED:
		k.handleTaskFailed(task, state, status)
	case mesos.TaskState_TASK_KILLED:
		k.handleTaskKilled(task, state, status)
	case mesos.TaskState_TASK_LOST:
		k.handleTaskLost(task, state, status)
	default:
		log.Warning("unhandled task status update: %+v", status)
	}
	return task, state
}

func (k *inMemoryRegistry) handleTaskStaging(task *T, state StateType, status *mesos.TaskStatus) {
	if status.GetSource() != mesos.TaskStatus_SOURCE_MASTER {
		log.Errorf("received STAGING for task %v with unexpected source: %v",
			status.GetTaskId().GetValue(), status.GetSource())
	}
}

func (k *inMemoryRegistry) handleTaskStarting(task *T, state StateType, status *mesos.TaskStatus) {
	// we expect to receive this when a launched task is finally "bound"
	// via the API server. however, there's nothing specific for us to do
	// here.
	switch state {
	case StatePending:
		task.UpdatedTime = time.Now()
		if !task.Has(Bound) {
			task.Set(Bound)
			task.bindTime = task.UpdatedTime
			timeToBind := task.bindTime.Sub(task.launchTime)
			metrics.BindLatency.Observe(metrics.InMicroseconds(timeToBind))
		}
	default:
		log.Warningf("Ignore status TASK_STARTING because the the task is not pending")
	}
}

func (k *inMemoryRegistry) handleTaskRunning(task *T, state StateType, status *mesos.TaskStatus) {
	switch state {
	case StatePending:
		task.UpdatedTime = time.Now()
		log.Infof("Received running status for pending task: %+v", status)
		fillRunningPodInfo(task, status)
		task.State = StateRunning
	case StateRunning:
		task.UpdatedTime = time.Now()
		log.V(2).Info("Ignore status TASK_RUNNING because the the task is already running")
	case StateFinished:
		log.Warningf("Ignore status TASK_RUNNING because the the task is already finished")
	default:
		log.Warningf("Ignore status TASK_RUNNING (%+v) because the the task is discarded", status.GetTaskId())
	}
}

func ParsePodStatusResult(taskStatus *mesos.TaskStatus) (result api.PodStatusResult, err error) {
	if taskStatus.Data != nil {
		err = json.Unmarshal(taskStatus.Data, &result)
	} else {
		err = fmt.Errorf("missing TaskStatus.Data")
	}
	return
}

func fillRunningPodInfo(task *T, taskStatus *mesos.TaskStatus) {
	if taskStatus.GetReason() == mesos.TaskStatus_REASON_RECONCILIATION && taskStatus.GetSource() == mesos.TaskStatus_SOURCE_MASTER {
		// there is no data..
		return
	}
	//TODO(jdef) determine the usefullness of this information (if any)
	if result, err := ParsePodStatusResult(taskStatus); err != nil {
		log.Errorf("invalid TaskStatus.Data for task '%v': %v", task.ID, err)
	} else {
		task.Pod.Status = result.Status
		log.Infof("received pod status for task %v: %+v", task.ID, task.Pod.Status)
	}
}

func (k *inMemoryRegistry) handleTaskFinished(task *T, state StateType, status *mesos.TaskStatus) {
	switch state {
	case StatePending:
		panic("Pending task finished, this couldn't happen")
	case StateRunning:
		log.V(2).Infof("received finished status for running task: %+v", status)
		delete(k.podToTask, task.podKey)
		task.State = StateFinished
		task.UpdatedTime = time.Now()
		k.tasksFinished = k.recordFinishedTask(task.ID)
	case StateFinished:
		log.Warningf("Ignore status TASK_FINISHED because the the task is already finished")
	default:
		log.Warningf("Ignore status TASK_FINISHED because the the task is not running")
	}
}

// record that a task has finished.
// older record are expunged one at a time once the historical ring buffer is saturated.
// assumes caller is holding state lock.
func (k *inMemoryRegistry) recordFinishedTask(taskId string) *ring.Ring {
	slot := k.tasksFinished.Next()
	if slot.Value != nil {
		// garbage collect older finished task from the registry
		gctaskId := slot.Value.(string)
		if gctask, found := k.taskRegistry[gctaskId]; found && gctask.State == StateFinished {
			delete(k.taskRegistry, gctaskId)
		}
	}
	slot.Value = taskId
	return slot
}

func (k *inMemoryRegistry) handleTaskFailed(task *T, state StateType, status *mesos.TaskStatus) {
	log.Errorf("task failed: %+v", status)
	switch state {
	case StatePending:
		delete(k.taskRegistry, task.ID)
		delete(k.podToTask, task.podKey)
	case StateRunning:
		delete(k.taskRegistry, task.ID)
		delete(k.podToTask, task.podKey)
	}
}

func (k *inMemoryRegistry) handleTaskKilled(task *T, state StateType, status *mesos.TaskStatus) {
	defer func() {
		msg := fmt.Sprintf("task killed: %+v, task %+v", status, task)
		if task != nil && task.Has(Deleted) {
			// we were expecting this, nothing out of the ordinary
			log.V(2).Infoln(msg)
		} else {
			log.Errorln(msg)
		}
	}()
	switch state {
	case StatePending, StateRunning:
		delete(k.taskRegistry, task.ID)
		delete(k.podToTask, task.podKey)
	}
}

func (k *inMemoryRegistry) handleTaskLost(task *T, state StateType, status *mesos.TaskStatus) {
	log.Warningf("task lost: %+v", status)
	switch state {
	case StateRunning, StatePending:
		delete(k.taskRegistry, task.ID)
		delete(k.podToTask, task.podKey)
	}
}
<|endoftext|>


<commit_before>package serviceaccounttoken

import (
	"context"
	"fmt"
	"time"

	"github.com/sirupsen/logrus"
	v1 "k8s.io/api/core/v1"
	apierror "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
)

// CreateSecretForServiceAccount creates a service-account-token Secret for the provided Service Account.
// If the secret already exists, the existing one is returned.
func CreateSecretForServiceAccount(ctx context.Context, clientSet kubernetes.Interface, sa *v1.ServiceAccount) (*v1.Secret, error) {
	secretName := ServiceAccountSecretName(sa)
	secretClient := clientSet.CoreV1().Secrets(sa.Namespace)
	secret, err := secretClient.Get(ctx, secretName, metav1.GetOptions{})
	if err != nil {
		if !apierror.IsNotFound(err) {
			return nil, err
		}
		sc := SecretTemplate(sa)
		secret, err = secretClient.Create(ctx, sc, metav1.CreateOptions{})
		if err != nil {
			if !apierror.IsAlreadyExists(err) {
				return nil, err
			}
			secret, err = secretClient.Get(ctx, secretName, metav1.GetOptions{})
			if err != nil {
				return nil, err
			}
		}
	}
	if len(secret.Data[v1.ServiceAccountTokenKey]) > 0 {
		return secret, nil
	}
	logrus.Infof("createSecretForServiceAccount: waiting for secret [%s] to be populated with token", secretName)
	for {
		if len(secret.Data[v1.ServiceAccountTokenKey]) > 0 {
			return secret, nil
		}
		time.Sleep(2 * time.Second)
		secret, err = secretClient.Get(ctx, secretName, metav1.GetOptions{})
		if err != nil {
			return nil, err
		}
	}
}

// SecretTemplate generate a template of service-account-token Secret for the provided Service Account.
func SecretTemplate(sa *v1.ServiceAccount) *v1.Secret {
	return &v1.Secret{
		ObjectMeta: metav1.ObjectMeta{
			Name:      ServiceAccountSecretName(sa),
			Namespace: sa.Namespace,
			OwnerReferences: []metav1.OwnerReference{
				{
					APIVersion: "v1",
					Kind:       "ServiceAccount",
					Name:       sa.Name,
					UID:        sa.UID,
				},
			},
			Annotations: map[string]string{
				"kubernetes.io/service-account.name": sa.Name,
			},
		},
		Type: v1.SecretTypeServiceAccountToken,
	}

}

// ServiceAccountSecretName returns the secret name for the given Service Account.
func ServiceAccountSecretName(sa *v1.ServiceAccount) string {
	return fmt.Sprintf("%s-token", sa.Name)
}
<commit_msg>Use name.SafeConcatName to ensure secret token name isn't longer than 63 characters<commit_after>package serviceaccounttoken

import (
	"context"
	"time"

	"github.com/rancher/wrangler/pkg/name"
	"github.com/sirupsen/logrus"
	v1 "k8s.io/api/core/v1"
	apierror "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
)

// CreateSecretForServiceAccount creates a service-account-token Secret for the provided Service Account.
// If the secret already exists, the existing one is returned.
func CreateSecretForServiceAccount(ctx context.Context, clientSet kubernetes.Interface, sa *v1.ServiceAccount) (*v1.Secret, error) {
	secretName := ServiceAccountSecretName(sa)
	secretClient := clientSet.CoreV1().Secrets(sa.Namespace)
	secret, err := secretClient.Get(ctx, secretName, metav1.GetOptions{})
	if err != nil {
		if !apierror.IsNotFound(err) {
			return nil, err
		}
		sc := SecretTemplate(sa)
		secret, err = secretClient.Create(ctx, sc, metav1.CreateOptions{})
		if err != nil {
			if !apierror.IsAlreadyExists(err) {
				return nil, err
			}
			secret, err = secretClient.Get(ctx, secretName, metav1.GetOptions{})
			if err != nil {
				return nil, err
			}
		}
	}
	if len(secret.Data[v1.ServiceAccountTokenKey]) > 0 {
		return secret, nil
	}
	logrus.Infof("createSecretForServiceAccount: waiting for secret [%s] to be populated with token", secretName)
	for {
		if len(secret.Data[v1.ServiceAccountTokenKey]) > 0 {
			return secret, nil
		}
		time.Sleep(2 * time.Second)
		secret, err = secretClient.Get(ctx, secretName, metav1.GetOptions{})
		if err != nil {
			return nil, err
		}
	}
}

// SecretTemplate generate a template of service-account-token Secret for the provided Service Account.
func SecretTemplate(sa *v1.ServiceAccount) *v1.Secret {
	return &v1.Secret{
		ObjectMeta: metav1.ObjectMeta{
			Name:      ServiceAccountSecretName(sa),
			Namespace: sa.Namespace,
			OwnerReferences: []metav1.OwnerReference{
				{
					APIVersion: "v1",
					Kind:       "ServiceAccount",
					Name:       sa.Name,
					UID:        sa.UID,
				},
			},
			Annotations: map[string]string{
				"kubernetes.io/service-account.name": sa.Name,
			},
		},
		Type: v1.SecretTypeServiceAccountToken,
	}

}

// ServiceAccountSecretName returns the secret name for the given Service Account.
func ServiceAccountSecretName(sa *v1.ServiceAccount) string {
	return name.SafeConcatName(sa.Name, "token")
}
<|endoftext|>


<commit_before>package prometheus

import (
	"context"
	"fmt"
	"regexp"
	"strings"
	"time"

	"net/http"

	"github.com/grafana/grafana/pkg/components/null"
	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/tsdb"
	"github.com/prometheus/client_golang/api/prometheus"
	pmodel "github.com/prometheus/common/model"
)

type PrometheusExecutor struct {
	*models.DataSource
	Transport *http.Transport
}

func NewPrometheusExecutor(dsInfo *models.DataSource) (tsdb.Executor, error) {
	transport, err := dsInfo.GetHttpTransport()
	if err != nil {
		return nil, err
	}

	return &PrometheusExecutor{
		DataSource: dsInfo,
		Transport:  transport,
	}, nil
}

var (
	plog         log.Logger
	legendFormat *regexp.Regexp
)

func init() {
	plog = log.New("tsdb.prometheus")
	tsdb.RegisterExecutor("prometheus", NewPrometheusExecutor)
	legendFormat = regexp.MustCompile(`\{\{\s*(.+?)\s*\}\}`)
}

func (e *PrometheusExecutor) getClient() (prometheus.QueryAPI, error) {
	cfg := prometheus.Config{
		Address:   e.DataSource.Url,
		Transport: e.Transport,
	}

	client, err := prometheus.New(cfg)
	if err != nil {
		return nil, err
	}

	return prometheus.NewQueryAPI(client), nil
}

func (e *PrometheusExecutor) Execute(ctx context.Context, queries tsdb.QuerySlice, queryContext *tsdb.QueryContext) *tsdb.BatchResult {
	result := &tsdb.BatchResult{}

	client, err := e.getClient()
	if err != nil {
		return result.WithError(err)
	}

	query, err := parseQuery(queries, queryContext)
	if err != nil {
		return result.WithError(err)
	}

	timeRange := prometheus.Range{
		Start: query.Start,
		End:   query.End,
		Step:  query.Step,
	}

	value, err := client.QueryRange(ctx, query.Expr, timeRange)

	if err != nil {
		return result.WithError(err)
	}

	queryResult, err := parseResponse(value, query)
	if err != nil {
		return result.WithError(err)
	}
	result.QueryResults = queryResult
	return result
}

func formatLegend(metric pmodel.Metric, query *PrometheusQuery) string {
	if query.LegendFormat == "" {
		return metric.String()
	}

	result := legendFormat.ReplaceAllFunc([]byte(query.LegendFormat), func(in []byte) []byte {
		labelName := strings.Replace(string(in), "{{", "", 1)
		labelName = strings.Replace(labelName, "}}", "", 1)
		labelName = strings.TrimSpace(labelName)
		if val, exists := metric[pmodel.LabelName(labelName)]; exists {
			return []byte(val)
		}

		return in
	})

	return string(result)
}

func parseQuery(queries tsdb.QuerySlice, queryContext *tsdb.QueryContext) (*PrometheusQuery, error) {
	queryModel := queries[0]

	expr, err := queryModel.Model.Get("expr").String()
	if err != nil {
		return nil, err
	}

	step, err := queryModel.Model.Get("step").Int64()
	if err != nil {
		return nil, err
	}

	format := queryModel.Model.Get("legendFormat").MustString("")

	start, err := queryContext.TimeRange.ParseFrom()
	if err != nil {
		return nil, err
	}

	end, err := queryContext.TimeRange.ParseTo()
	if err != nil {
		return nil, err
	}

	return &PrometheusQuery{
		Expr:         expr,
		Step:         time.Second * time.Duration(step),
		LegendFormat: format,
		Start:        start,
		End:          end,
	}, nil
}

func parseResponse(value pmodel.Value, query *PrometheusQuery) (map[string]*tsdb.QueryResult, error) {
	queryResults := make(map[string]*tsdb.QueryResult)
	queryRes := tsdb.NewQueryResult()

	data, ok := value.(pmodel.Matrix)
	if !ok {
		return queryResults, fmt.Errorf("Unsupported result format: %s", value.Type().String())
	}

	for _, v := range data {
		series := tsdb.TimeSeries{
			Name: formatLegend(v.Metric, query),
		}

		for _, k := range v.Values {
			series.Points = append(series.Points, tsdb.NewTimePoint(null.FloatFrom(float64(k.Value)), float64(k.Timestamp.Unix()*1000)))
		}

		queryRes.Series = append(queryRes.Series, &series)
	}

	queryResults["A"] = queryRes
	return queryResults, nil
}
<commit_msg>feat(alerting): transform labels into tags for prometheus tsdb<commit_after>package prometheus

import (
	"context"
	"fmt"
	"regexp"
	"strings"
	"time"

	"net/http"

	"github.com/grafana/grafana/pkg/components/null"
	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/tsdb"
	"github.com/prometheus/client_golang/api/prometheus"
	pmodel "github.com/prometheus/common/model"
)

type PrometheusExecutor struct {
	*models.DataSource
	Transport *http.Transport
}

func NewPrometheusExecutor(dsInfo *models.DataSource) (tsdb.Executor, error) {
	transport, err := dsInfo.GetHttpTransport()
	if err != nil {
		return nil, err
	}

	return &PrometheusExecutor{
		DataSource: dsInfo,
		Transport:  transport,
	}, nil
}

var (
	plog         log.Logger
	legendFormat *regexp.Regexp
)

func init() {
	plog = log.New("tsdb.prometheus")
	tsdb.RegisterExecutor("prometheus", NewPrometheusExecutor)
	legendFormat = regexp.MustCompile(`\{\{\s*(.+?)\s*\}\}`)
}

func (e *PrometheusExecutor) getClient() (prometheus.QueryAPI, error) {
	cfg := prometheus.Config{
		Address:   e.DataSource.Url,
		Transport: e.Transport,
	}

	client, err := prometheus.New(cfg)
	if err != nil {
		return nil, err
	}

	return prometheus.NewQueryAPI(client), nil
}

func (e *PrometheusExecutor) Execute(ctx context.Context, queries tsdb.QuerySlice, queryContext *tsdb.QueryContext) *tsdb.BatchResult {
	result := &tsdb.BatchResult{}

	client, err := e.getClient()
	if err != nil {
		return result.WithError(err)
	}

	query, err := parseQuery(queries, queryContext)
	if err != nil {
		return result.WithError(err)
	}

	timeRange := prometheus.Range{
		Start: query.Start,
		End:   query.End,
		Step:  query.Step,
	}

	value, err := client.QueryRange(ctx, query.Expr, timeRange)

	if err != nil {
		return result.WithError(err)
	}

	queryResult, err := parseResponse(value, query)
	if err != nil {
		return result.WithError(err)
	}
	result.QueryResults = queryResult
	return result
}

func formatLegend(metric pmodel.Metric, query *PrometheusQuery) string {
	if query.LegendFormat == "" {
		return metric.String()
	}

	result := legendFormat.ReplaceAllFunc([]byte(query.LegendFormat), func(in []byte) []byte {
		labelName := strings.Replace(string(in), "{{", "", 1)
		labelName = strings.Replace(labelName, "}}", "", 1)
		labelName = strings.TrimSpace(labelName)
		if val, exists := metric[pmodel.LabelName(labelName)]; exists {
			return []byte(val)
		}

		return in
	})

	return string(result)
}

func parseQuery(queries tsdb.QuerySlice, queryContext *tsdb.QueryContext) (*PrometheusQuery, error) {
	queryModel := queries[0]

	expr, err := queryModel.Model.Get("expr").String()
	if err != nil {
		return nil, err
	}

	step, err := queryModel.Model.Get("step").Int64()
	if err != nil {
		return nil, err
	}

	format := queryModel.Model.Get("legendFormat").MustString("")

	start, err := queryContext.TimeRange.ParseFrom()
	if err != nil {
		return nil, err
	}

	end, err := queryContext.TimeRange.ParseTo()
	if err != nil {
		return nil, err
	}

	return &PrometheusQuery{
		Expr:         expr,
		Step:         time.Second * time.Duration(step),
		LegendFormat: format,
		Start:        start,
		End:          end,
	}, nil
}

func parseResponse(value pmodel.Value, query *PrometheusQuery) (map[string]*tsdb.QueryResult, error) {
	queryResults := make(map[string]*tsdb.QueryResult)
	queryRes := tsdb.NewQueryResult()

	data, ok := value.(pmodel.Matrix)
	if !ok {
		return queryResults, fmt.Errorf("Unsupported result format: %s", value.Type().String())
	}

	for _, v := range data {
		series := tsdb.TimeSeries{
			Name: formatLegend(v.Metric, query),
			Tags: map[string]string{},
		}

		for k, v := range v.Metric {
			series.Tags[string(k)] = string(v)
		}

		for _, k := range v.Values {
			series.Points = append(series.Points, tsdb.NewTimePoint(null.FloatFrom(float64(k.Value)), float64(k.Timestamp.Unix()*1000)))
		}

		queryRes.Series = append(queryRes.Series, &series)
	}

	queryResults["A"] = queryRes
	return queryResults, nil
}
<|endoftext|>


<commit_before>package products

import (
	"time"
)

// Product represents a software product within the system for identification
// across multiple sources
type Product struct {
	ID         int           `json:"id" xml:"id"`
	Name       string        `json:"name" xml:"name"`
	Org        string        `json:"org" xml:"org"`
	Version    string        `json:"version" xml:"version"`
	Up         string        `json:"up" xml:"up"`
	Edition    string        `json:"edition" xml:"edition"`
	Aliases    interface{}   `json:"aliases" xml:"aliases"`
	CreatedAt  time.Time     `json:"created_at" xml:"created_at"`
	UpdatedAt  time.Time     `json:"updated_at" xml:"updated_at"`
	Title      string        `json:"title" xml:"title"`
	References []interface{} `json:"references" xml:"references"`
	Part       string        `json:"part" xml:"part"`
	Language   string        `json:"language" xml:"language"`
	ExternalID string        `json:"external_id" xml:"external_id"`
	Sources    []Source      `json:"source" xml:"source"`
}

// Source represents information about where the product data came from
type Source struct {
	ID           int       `json:"id" xml:"id"`
	Name         string    `json:"name" xml:"name"`
	Description  string    `json:"description" xml:"description"`
	CreatedAt    time.Time `json:"created_at" xml:"created_at"`
	UpdatedAt    time.Time `json:"updated_at" xml:"updated_at"`
	Attribution  string    `json:"attribution" xml:"attribution"`
	License      string    `json:"license" xml:"license"`
	CopyrightURL string    `json:"copyright_url" xml:"copyright_url"`
}

// SoftwareEntity represents information about a product as well as
// other info, like Git repository, committer counts, etc
type SoftwareEntity struct {
	Product    *Product              `json:"product,omitempty" xml:"product"`
	Github     *Github               `json:"github,omitempty" xml:"github,omitempty"`
	Package    *Package              `json:"package,omitempty" xml:"package,omitempty"`
	Confidence float64              `json:"mean_score" xml:"mean_score"`
	Scores     []ProductSearchScore `json:"scores,omitempty" xml:"scores"`
}

// ProductSearchScore represents the TF;IDF score for a given search result
// and a given search term
type ProductSearchScore struct {
	Term  string  `json:"term" xml:"term"`
	Score float64 `json:"score" xml:"score"`
}

// Github represents information from Github about a given repository
type Github struct {
	URI            string `json:"uri" xml:"uri"`
	CommitterCount uint   `json:"committer_count" xml:"committer_count"`
}

// Package represents information about a package from one of
// our supported package management systems like pypi, npm or rubygems
type Package struct {
	Name    string `json:"name" xml:"name"`
	Version string `json:"version" xml:"version"`
	Type    string `json:"type" xml:"type"`
}

// ProductSearchQuery collects all the various searching options that
// the productSearchEndpoint supports for use in a POST request
type ProductSearchQuery struct {
	SearchType        string   `json:"search_type" xml:"search_type"`
	SearchStrategy    string   `json:"search_strategy" xml:"search_strategy"`
	ProductIdentifier string   `json:"product_identifier" xml:"product_identifier"`
	Version           string   `json:"version" xml:"version"`
	Vendor            string   `json:"vendor" xml:"vendor"`
	Terms             []string `json:"terms" xml:"terms"`
}

// IsValid checks some of the constraints on the ProductSearchQuery to
// help the programmer determine if productSearchEndpoint will accept it
func (p *ProductSearchQuery) IsValid() bool {
	if len(p.SearchStrategy) > 0 {
		if p.SearchType == "concatenated" || p.SearchType == "deconcatenated" {
			return true
		}
	}
	return false
}
<commit_msg>also changing mean_score to confidence in json/xml tags<commit_after>package products

import (
	"time"
)

// Product represents a software product within the system for identification
// across multiple sources
type Product struct {
	ID         int           `json:"id" xml:"id"`
	Name       string        `json:"name" xml:"name"`
	Org        string        `json:"org" xml:"org"`
	Version    string        `json:"version" xml:"version"`
	Up         string        `json:"up" xml:"up"`
	Edition    string        `json:"edition" xml:"edition"`
	Aliases    interface{}   `json:"aliases" xml:"aliases"`
	CreatedAt  time.Time     `json:"created_at" xml:"created_at"`
	UpdatedAt  time.Time     `json:"updated_at" xml:"updated_at"`
	Title      string        `json:"title" xml:"title"`
	References []interface{} `json:"references" xml:"references"`
	Part       string        `json:"part" xml:"part"`
	Language   string        `json:"language" xml:"language"`
	ExternalID string        `json:"external_id" xml:"external_id"`
	Sources    []Source      `json:"source" xml:"source"`
}

// Source represents information about where the product data came from
type Source struct {
	ID           int       `json:"id" xml:"id"`
	Name         string    `json:"name" xml:"name"`
	Description  string    `json:"description" xml:"description"`
	CreatedAt    time.Time `json:"created_at" xml:"created_at"`
	UpdatedAt    time.Time `json:"updated_at" xml:"updated_at"`
	Attribution  string    `json:"attribution" xml:"attribution"`
	License      string    `json:"license" xml:"license"`
	CopyrightURL string    `json:"copyright_url" xml:"copyright_url"`
}

// SoftwareEntity represents information about a product as well as
// other info, like Git repository, committer counts, etc
type SoftwareEntity struct {
	Product    *Product             `json:"product,omitempty" xml:"product"`
	Github     *Github              `json:"github,omitempty" xml:"github,omitempty"`
	Package    *Package             `json:"package,omitempty" xml:"package,omitempty"`
	Confidence float64              `json:"confidence" xml:"confidence"`
	Scores     []ProductSearchScore `json:"scores,omitempty" xml:"scores"`
}

// ProductSearchScore represents the TF;IDF score for a given search result
// and a given search term
type ProductSearchScore struct {
	Term  string  `json:"term" xml:"term"`
	Score float64 `json:"score" xml:"score"`
}

// Github represents information from Github about a given repository
type Github struct {
	URI            string `json:"uri" xml:"uri"`
	CommitterCount uint   `json:"committer_count" xml:"committer_count"`
}

// Package represents information about a package from one of
// our supported package management systems like pypi, npm or rubygems
type Package struct {
	Name    string `json:"name" xml:"name"`
	Version string `json:"version" xml:"version"`
	Type    string `json:"type" xml:"type"`
}

// ProductSearchQuery collects all the various searching options that
// the productSearchEndpoint supports for use in a POST request
type ProductSearchQuery struct {
	SearchType        string   `json:"search_type" xml:"search_type"`
	SearchStrategy    string   `json:"search_strategy" xml:"search_strategy"`
	ProductIdentifier string   `json:"product_identifier" xml:"product_identifier"`
	Version           string   `json:"version" xml:"version"`
	Vendor            string   `json:"vendor" xml:"vendor"`
	Terms             []string `json:"terms" xml:"terms"`
}

// IsValid checks some of the constraints on the ProductSearchQuery to
// help the programmer determine if productSearchEndpoint will accept it
func (p *ProductSearchQuery) IsValid() bool {
	if len(p.SearchStrategy) > 0 {
		if p.SearchType == "concatenated" || p.SearchType == "deconcatenated" {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>// Copyright 2015 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// This file does serialization of programs for executor binary.
// The format aims at simple parsing: binary and irreversible.

package prog

import (
	"fmt"

	"github.com/google/syzkaller/sys"
)

const (
	ExecInstrEOF = ^uintptr(iota)
	ExecInstrCopyin
	ExecInstrCopyout
)

const (
	ExecArgConst = uintptr(iota)
	ExecArgResult
	ExecArgData
)

const (
	ptrSize    = 8
	pageSize   = 4 << 10
	dataOffset = 512 << 20
)

func (p *Prog) SerializeForExec(pid int) []byte {
	if err := p.validate(); err != nil {
		panic(fmt.Errorf("serializing invalid program: %v", err))
	}
	var instrSeq uintptr
	w := &execContext{args: make(map[*Arg]*argInfo)}
	for _, c := range p.Calls {
		// Calculate arg offsets within structs.
		foreachArg(c, func(arg, base *Arg, _ *[]*Arg) {
			if base == nil || arg.Kind == ArgGroup || arg.Kind == ArgUnion {
				return
			}
			if w.args[base] == nil {
				w.args[base] = &argInfo{}
			}
			w.args[arg] = &argInfo{Offset: w.args[base].CurSize}
			if arg.Type.BitfieldLength() == 0 || arg.Type.BitfieldLast() {
				w.args[base].CurSize += arg.Size()
			}
		})
		// Generate copyin instructions that fill in data into pointer arguments.
		foreachArg(c, func(arg, _ *Arg, _ *[]*Arg) {
			if arg.Kind == ArgPointer && arg.Res != nil {
				var rec func(*Arg)
				rec = func(arg1 *Arg) {
					if arg1.Kind == ArgGroup {
						for _, arg2 := range arg1.Inner {
							rec(arg2)
						}
						return
					}
					if arg1.Kind == ArgUnion {
						rec(arg1.Option)
						return
					}
					if sys.IsPad(arg1.Type) {
						return
					}
					if arg1.Kind == ArgData && len(arg1.Data) == 0 {
						return
					}
					if arg1.Type.Dir() != sys.DirOut {
						w.write(ExecInstrCopyin)
						w.write(physicalAddr(arg) + w.args[arg1].Offset)
						w.writeArg(arg1, pid)
						instrSeq++
					}
				}
				rec(arg.Res)
			}
		})
		// Generate the call itself.
		w.write(uintptr(c.Meta.ID))
		w.write(uintptr(len(c.Args)))
		for _, arg := range c.Args {
			w.writeArg(arg, pid)
		}
		w.args[c.Ret] = &argInfo{Idx: instrSeq}
		instrSeq++
		// Generate copyout instructions that persist interesting return values.
		foreachArg(c, func(arg, base *Arg, _ *[]*Arg) {
			if len(arg.Uses) == 0 {
				return
			}
			switch arg.Kind {
			case ArgReturn:
				// Idx is already assigned above.
			case ArgConst, ArgResult:
				// Create a separate copyout instruction that has own Idx.
				if base.Kind != ArgPointer {
					panic("arg base is not a pointer")
				}
				info := w.args[arg]
				info.Idx = instrSeq
				instrSeq++
				w.write(ExecInstrCopyout)
				w.write(physicalAddr(base) + info.Offset)
				w.write(arg.Size())
			default:
				panic("bad arg kind in copyout")
			}
		})
	}
	w.write(ExecInstrEOF)
	return w.buf
}

func physicalAddr(arg *Arg) uintptr {
	if arg.Kind != ArgPointer {
		panic("physicalAddr: bad arg kind")
	}
	addr := arg.AddrPage*pageSize + dataOffset
	if arg.AddrOffset >= 0 {
		addr += uintptr(arg.AddrOffset)
	} else {
		addr += pageSize - uintptr(-arg.AddrOffset)
	}
	return addr
}

type execContext struct {
	buf  []byte
	args map[*Arg]*argInfo
}

type argInfo struct {
	Offset  uintptr // from base pointer
	CurSize uintptr
	Idx     uintptr // instruction index
}

func (w *execContext) write(v uintptr) {
	w.buf = append(w.buf, byte(v>>0), byte(v>>8), byte(v>>16), byte(v>>24), byte(v>>32), byte(v>>40), byte(v>>48), byte(v>>56))
}

func (w *execContext) writeArg(arg *Arg, pid int) {
	switch arg.Kind {
	case ArgConst:
		w.write(ExecArgConst)
		w.write(arg.Size())
		w.write(arg.Value(pid))
		w.write(arg.Type.BitfieldOffset())
		w.write(arg.Type.BitfieldLength())
	case ArgResult:
		w.write(ExecArgResult)
		w.write(arg.Size())
		w.write(w.args[arg.Res].Idx)
		w.write(arg.OpDiv)
		w.write(arg.OpAdd)
	case ArgPointer:
		w.write(ExecArgConst)
		w.write(arg.Size())
		w.write(physicalAddr(arg))
		w.write(0) // bit field offset
		w.write(0) // bit field length
	case ArgPageSize:
		w.write(ExecArgConst)
		w.write(arg.Size())
		w.write(arg.AddrPage * pageSize)
		w.write(0) // bit field offset
		w.write(0) // bit field length
	case ArgData:
		w.write(ExecArgData)
		w.write(uintptr(len(arg.Data)))
		for i := 0; i < len(arg.Data); i += 8 {
			var v uintptr
			for j := 0; j < 8; j++ {
				if i+j >= len(arg.Data) {
					break
				}
				v |= uintptr(arg.Data[i+j]) << uint(j*8)
			}
			w.write(v)
		}
	default:
		panic("unknown arg type")
	}
}
<commit_msg>prog: fix union and struct offsets in SerializeForExec<commit_after>// Copyright 2015 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// This file does serialization of programs for executor binary.
// The format aims at simple parsing: binary and irreversible.

package prog

import (
	"fmt"

	"github.com/google/syzkaller/sys"
)

const (
	ExecInstrEOF = ^uintptr(iota)
	ExecInstrCopyin
	ExecInstrCopyout
)

const (
	ExecArgConst = uintptr(iota)
	ExecArgResult
	ExecArgData
)

const (
	ptrSize    = 8
	pageSize   = 4 << 10
	dataOffset = 512 << 20
)

func (p *Prog) SerializeForExec(pid int) []byte {
	if err := p.validate(); err != nil {
		panic(fmt.Errorf("serializing invalid program: %v", err))
	}
	var instrSeq uintptr
	w := &execContext{args: make(map[*Arg]*argInfo)}
	for _, c := range p.Calls {
		// Calculate arg offsets within structs.
		// Generate copyin instructions that fill in data into pointer arguments.
		foreachArg(c, func(arg, _ *Arg, _ *[]*Arg) {
			if arg.Kind == ArgPointer && arg.Res != nil {
				var rec func(*Arg, uintptr) uintptr
				rec = func(arg1 *Arg, offset uintptr) uintptr {
					w.args[arg1] = &argInfo{Offset: offset}
					if arg1.Kind == ArgGroup {
						var totalSize uintptr
						for _, arg2 := range arg1.Inner {
							size := rec(arg2, offset)
							if arg2.Type.BitfieldLength() == 0 || arg2.Type.BitfieldLast() {
								offset += size
								totalSize += size
							}
						}
						if totalSize > arg1.Size() {
							panic(fmt.Sprintf("bad group arg size %v, should be <= %v for %+v", totalSize, arg1.Size(), arg1))
						}
						return arg1.Size()
					}
					if arg1.Kind == ArgUnion {
						size := rec(arg1.Option, offset)
						offset += size
						if size > arg1.Size() {
							panic(fmt.Sprintf("bad union arg size %v, should be <= %v for %+v", size, arg1.Size(), arg1))
						}
						return arg1.Size()
					}
					if !sys.IsPad(arg1.Type) &&
						!(arg1.Kind == ArgData && len(arg1.Data) == 0) &&
						arg1.Type.Dir() != sys.DirOut {
						w.write(ExecInstrCopyin)
						w.write(physicalAddr(arg) + offset)
						w.writeArg(arg1, pid)
						instrSeq++
					}
					return arg1.Size()
				}
				rec(arg.Res, 0)
			}
		})
		// Generate the call itself.
		w.write(uintptr(c.Meta.ID))
		w.write(uintptr(len(c.Args)))
		for _, arg := range c.Args {
			w.writeArg(arg, pid)
		}
		w.args[c.Ret] = &argInfo{Idx: instrSeq}
		instrSeq++
		// Generate copyout instructions that persist interesting return values.
		foreachArg(c, func(arg, base *Arg, _ *[]*Arg) {
			if len(arg.Uses) == 0 {
				return
			}
			switch arg.Kind {
			case ArgReturn:
				// Idx is already assigned above.
			case ArgConst, ArgResult:
				// Create a separate copyout instruction that has own Idx.
				if base.Kind != ArgPointer {
					panic("arg base is not a pointer")
				}
				info := w.args[arg]
				info.Idx = instrSeq
				instrSeq++
				w.write(ExecInstrCopyout)
				w.write(physicalAddr(base) + info.Offset)
				w.write(arg.Size())
			default:
				panic("bad arg kind in copyout")
			}
		})
	}
	w.write(ExecInstrEOF)
	return w.buf
}

func physicalAddr(arg *Arg) uintptr {
	if arg.Kind != ArgPointer {
		panic("physicalAddr: bad arg kind")
	}
	addr := arg.AddrPage*pageSize + dataOffset
	if arg.AddrOffset >= 0 {
		addr += uintptr(arg.AddrOffset)
	} else {
		addr += pageSize - uintptr(-arg.AddrOffset)
	}
	return addr
}

type execContext struct {
	buf  []byte
	args map[*Arg]*argInfo
}

type argInfo struct {
	Offset uintptr // from base pointer
	Idx    uintptr // instruction index
}

func (w *execContext) write(v uintptr) {
	w.buf = append(w.buf, byte(v>>0), byte(v>>8), byte(v>>16), byte(v>>24), byte(v>>32), byte(v>>40), byte(v>>48), byte(v>>56))
}

func (w *execContext) writeArg(arg *Arg, pid int) {
	switch arg.Kind {
	case ArgConst:
		w.write(ExecArgConst)
		w.write(arg.Size())
		w.write(arg.Value(pid))
		w.write(arg.Type.BitfieldOffset())
		w.write(arg.Type.BitfieldLength())
	case ArgResult:
		w.write(ExecArgResult)
		w.write(arg.Size())
		w.write(w.args[arg.Res].Idx)
		w.write(arg.OpDiv)
		w.write(arg.OpAdd)
	case ArgPointer:
		w.write(ExecArgConst)
		w.write(arg.Size())
		w.write(physicalAddr(arg))
		w.write(0) // bit field offset
		w.write(0) // bit field length
	case ArgPageSize:
		w.write(ExecArgConst)
		w.write(arg.Size())
		w.write(arg.AddrPage * pageSize)
		w.write(0) // bit field offset
		w.write(0) // bit field length
	case ArgData:
		w.write(ExecArgData)
		w.write(uintptr(len(arg.Data)))
		for i := 0; i < len(arg.Data); i += 8 {
			var v uintptr
			for j := 0; j < 8; j++ {
				if i+j >= len(arg.Data) {
					break
				}
				v |= uintptr(arg.Data[i+j]) << uint(j*8)
			}
			w.write(v)
		}
	default:
		panic("unknown arg type")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/hex"
	"fmt"
	"strings"
	"sync"

	"github.com/pkg/term"
	"go.bug.st/serial.v1"
)

const NODE_ID = "3D"

var (
	port serial.Port

	remote   = false
	sendLock sync.Mutex // protects sendBuf between goroutine and main
	sendBuf  string
)

func main() {
	var e error
	device := "/dev/cu.usbmodem32212431"
	port, e = serial.Open(device, &serial.Mode{BaudRate: 115200})
	if e != nil {
		panic(e)
	}
	defer port.Close()
	fmt.Println("[Connected]")

	go handleSerialInput()

	handleConsoleInput()
}

// handleConsoleInput handles key input, either to send directly or over RF.
func handleConsoleInput() {
	t, e := term.Open("/dev/tty")
	if e != nil {
		panic(e)
	}
	defer t.Restore()
	term.RawMode(t)

	key := make([]byte, 1)
	for {
		n, e := t.Read(key)
		if e != nil || n == 0 {
			break
		}

		switch key[0] {
		case 0x04: // Ctrl-D, quit
			return
		case 0x12: // Ctrl-R, enter remote mode
			remote = true
			sendBuf = ""
			fmt.Print("\r\n[Remote]\r\n")
		case 0x0C: // Ctrl-L, exit remote mode, back to local
			remote = false
			fmt.Print("\r\n[Local]\r\n")
		default:
			if remote {
				sendLock.Lock()
				sendBuf += string(key)
				sendLock.Unlock()
			} else {
				port.Write(key)
			}
		}
	}
}

// handleSerialInput deals with all incoming serial data.
func handleSerialInput() {
	recvBuf := ""
	for {
		data := make([]byte, 250)
		n, e := port.Read(data)
		if e != nil {
			break
		}
		s := string(data[:n])
		if remote {
			recvBuf += s
			sendLock.Lock()
			if parsePacketMsg(recvBuf) {
				recvBuf = ""
			}
			sendLock.Unlock()
		} else {
			s = strings.Replace(s, "\n", "\r\n", -1)
			fmt.Print(s)
		}
	}
}

// parsePacketMsg recognises incoming RF69 data packets
func parsePacketMsg(s string) bool {
	if pos := strings.LastIndex(s, "RF69 "); pos >= 0 {
		s = s[pos:]
		if strings.Contains(s, "\n") {
			// RF69 21EE06AB01005EC0010A 8111F209D017994EB780\n
			f := strings.Fields(s)
			if len(f) == 3 && len(f[1]) == 20 && f[1][16:18] == NODE_ID {
				if b, e := hex.DecodeString(f[2]); e == nil {
					processPacket(b)
					return true
				}
			}
			fmt.Printf("%q ?\r\n", s)
		}
	}
	return false
}

// processPacket responds to each incoming packet by sending an ACK packet
func processPacket(packet []byte) {
	fmt.Print(packet, "\r\n")
}
<commit_msg>implement packet reply logic<commit_after>package main

import (
	"encoding/hex"
	"fmt"
	"strings"
	"sync"

	"github.com/pkg/term"
	"go.bug.st/serial.v1"
)

const NODE_ID = "3D"

var (
	port serial.Port

	remote   = false
	sendLock sync.Mutex // protects sendBuf between goroutine and main
	sendBuf  string
)

func main() {
	var e error
	device := "/dev/cu.usbmodem32212431"
	port, e = serial.Open(device, &serial.Mode{BaudRate: 115200})
	if e != nil {
		panic(e)
	}
	defer port.Close()
	fmt.Println("[Connected]")

	go handleSerialInput()

	handleConsoleInput()
}

// handleConsoleInput handles key input, either to send directly or over RF.
func handleConsoleInput() {
	t, e := term.Open("/dev/tty")
	if e != nil {
		panic(e)
	}
	defer t.Restore()
	term.RawMode(t)

	key := make([]byte, 1)
	for {
		n, e := t.Read(key)
		if e != nil || n == 0 {
			break
		}

		switch key[0] {
		case 0x04: // Ctrl-D, quit
			return
		case 0x12: // Ctrl-R, enter remote mode
			remote = true
			sendBuf = ""
			fmt.Print("\r\n[Remote]\r\n")
		case 0x0C: // Ctrl-L, exit remote mode, back to local
			remote = false
			fmt.Print("\r\n[Local]\r\n")
		default:
			if remote {
				sendLock.Lock()
				sendBuf += string(key)
				sendLock.Unlock()
			} else {
				port.Write(key)
			}
		}
	}
}

// handleSerialInput deals with all incoming serial data.
func handleSerialInput() {
	recvBuf := ""
	for {
		data := make([]byte, 250)
		n, e := port.Read(data)
		if e != nil {
			break
		}
		s := string(data[:n])
		if remote {
			recvBuf += s
			sendLock.Lock()
			if parsePacketMsg(recvBuf) {
				recvBuf = ""
			}
			sendLock.Unlock()
		} else {
			display(s)
		}
	}
}

// display shows received data while inserting a CR before each LF
func display(s string) {
	s = strings.Replace(s, "\n", "\r\n", -1)
	fmt.Print(s)
}

// parsePacketMsg recognises incoming RF69 data packets
func parsePacketMsg(s string) bool {
	if pos := strings.LastIndex(s, "RF69 "); pos >= 0 {
		s = s[pos:]
		if strings.Contains(s, "\n") {
			// RF69 21EE06AB01005EC0010A 8111F209D017994EB780\n
			f := strings.Fields(s)
			if len(f) == 3 && len(f[1]) == 20 && f[1][16:18] == NODE_ID {
				if b, e := hex.DecodeString(f[2]); e == nil {
					processPacket(b)
					return true
				}
			}
			fmt.Printf("%q ?\r\n", s)
		}
	}
	return false
}

// processPacket responds to each incoming packet by sending an ACK packet
func processPacket(p []byte) {
	display(string(p[1:]))
	msg := append([]byte{0}, []byte(sendBuf)...)
	if len(msg) > 66 {
		msg = msg[:66]
	}
	sendBuf = sendBuf[len(msg)-1:]
	fmt.Print(msg, "\r\n")
}
<|endoftext|>


<commit_before>package terminal

import (
	"os"
	"regexp"

	"github.com/fatih/color"
	"golang.org/x/crypto/ssh/terminal"
)

const (
	red     color.Attribute = color.FgRed
	green                   = color.FgGreen
	yellow                  = color.FgYellow
	magenta                 = color.FgMagenta
	cyan                    = color.FgCyan
	grey                    = color.FgWhite
	white                   = color.FgHiWhite
)

var (
	colorize               func(message string, textColor color.Attribute, bold int) string
	TerminalSupportsColors = isTerminal()
	UserAskedForColors     = ""
)

func init() {
	InitColorSupport()
}

func InitColorSupport() {
	if colorsEnabled() {
		colorize = func(message string, textColor color.Attribute, bold int) string {
			colorPrinter := color.New(textColor)
			if bold == 1 {
				colorPrinter = colorPrinter.Add(color.Bold)
			}
			f := colorPrinter.SprintFunc()
			return f(message)
		}
	} else {
		colorize = func(message string, _ color.Attribute, _ int) string {
			return message
		}
	}
}

func colorsEnabled() bool {
	if os.Getenv("CF_COLOR") == "true" {
		return true
	}

	if os.Getenv("CF_COLOR") == "false" {
		return false
	}

	if UserAskedForColors == "true" {
		return true
	}

	return UserAskedForColors != "false" && TerminalSupportsColors
}

func Colorize(message string, textColor color.Attribute) string {
	return colorize(message, textColor, 0)
}

func ColorizeBold(message string, textColor color.Attribute) string {
	return colorize(message, textColor, 1)
}

var decolorizerRegex = regexp.MustCompile(`\x1B\[([0-9]{1,2}(;[0-9]{1,2})?)?[m|K]`)

func Decolorize(message string) string {
	return string(decolorizerRegex.ReplaceAll([]byte(message), []byte("")))
}

func HeaderColor(message string) string {
	return ColorizeBold(message, white)
}

func CommandColor(message string) string {
	return ColorizeBold(message, yellow)
}

func StoppedColor(message string) string {
	return ColorizeBold(message, grey)
}

func AdvisoryColor(message string) string {
	return ColorizeBold(message, yellow)
}

func CrashedColor(message string) string {
	return ColorizeBold(message, red)
}

func FailureColor(message string) string {
	return ColorizeBold(message, red)
}

func SuccessColor(message string) string {
	return ColorizeBold(message, green)
}

func EntityNameColor(message string) string {
	return ColorizeBold(message, cyan)
}

func PromptColor(message string) string {
	return ColorizeBold(message, cyan)
}

func TableContentHeaderColor(message string) string {
	return ColorizeBold(message, cyan)
}

func WarningColor(message string) string {
	return ColorizeBold(message, magenta)
}

func LogStdoutColor(message string) string {
	return Colorize(message, white)
}

func LogStderrColor(message string) string {
	return Colorize(message, red)
}

func LogHealthHeaderColor(message string) string {
	return Colorize(message, grey)
}

func LogAppHeaderColor(message string) string {
	return ColorizeBold(message, yellow)
}

func LogSysHeaderColor(message string) string {
	return ColorizeBold(message, cyan)
}

func isTerminal() bool {
	return terminal.IsTerminal(int(os.Stdout.Fd()))
}
<commit_msg>Use Normal White instead of Bright White<commit_after>package terminal

import (
	"os"
	"regexp"

	"github.com/fatih/color"
	"golang.org/x/crypto/ssh/terminal"
)

const (
	red     color.Attribute = color.FgRed
	green                   = color.FgGreen
	yellow                  = color.FgYellow
	magenta                 = color.FgMagenta
	cyan                    = color.FgCyan
	grey                    = color.FgWhite
	white                   = color.FgWhite
)

var (
	colorize               func(message string, textColor color.Attribute, bold int) string
	TerminalSupportsColors = isTerminal()
	UserAskedForColors     = ""
)

func init() {
	InitColorSupport()
}

func InitColorSupport() {
	if colorsEnabled() {
		colorize = func(message string, textColor color.Attribute, bold int) string {
			colorPrinter := color.New(textColor)
			if bold == 1 {
				colorPrinter = colorPrinter.Add(color.Bold)
			}
			f := colorPrinter.SprintFunc()
			return f(message)
		}
	} else {
		colorize = func(message string, _ color.Attribute, _ int) string {
			return message
		}
	}
}

func colorsEnabled() bool {
	if os.Getenv("CF_COLOR") == "true" {
		return true
	}

	if os.Getenv("CF_COLOR") == "false" {
		return false
	}

	if UserAskedForColors == "true" {
		return true
	}

	return UserAskedForColors != "false" && TerminalSupportsColors
}

func Colorize(message string, textColor color.Attribute) string {
	return colorize(message, textColor, 0)
}

func ColorizeBold(message string, textColor color.Attribute) string {
	return colorize(message, textColor, 1)
}

var decolorizerRegex = regexp.MustCompile(`\x1B\[([0-9]{1,2}(;[0-9]{1,2})?)?[m|K]`)

func Decolorize(message string) string {
	return string(decolorizerRegex.ReplaceAll([]byte(message), []byte("")))
}

func HeaderColor(message string) string {
	return ColorizeBold(message, white)
}

func CommandColor(message string) string {
	return ColorizeBold(message, yellow)
}

func StoppedColor(message string) string {
	return ColorizeBold(message, grey)
}

func AdvisoryColor(message string) string {
	return ColorizeBold(message, yellow)
}

func CrashedColor(message string) string {
	return ColorizeBold(message, red)
}

func FailureColor(message string) string {
	return ColorizeBold(message, red)
}

func SuccessColor(message string) string {
	return ColorizeBold(message, green)
}

func EntityNameColor(message string) string {
	return ColorizeBold(message, cyan)
}

func PromptColor(message string) string {
	return ColorizeBold(message, cyan)
}

func TableContentHeaderColor(message string) string {
	return ColorizeBold(message, cyan)
}

func WarningColor(message string) string {
	return ColorizeBold(message, magenta)
}

func LogStdoutColor(message string) string {
	return Colorize(message, white)
}

func LogStderrColor(message string) string {
	return Colorize(message, red)
}

func LogHealthHeaderColor(message string) string {
	return Colorize(message, grey)
}

func LogAppHeaderColor(message string) string {
	return ColorizeBold(message, yellow)
}

func LogSysHeaderColor(message string) string {
	return ColorizeBold(message, cyan)
}

func isTerminal() bool {
	return terminal.IsTerminal(int(os.Stdout.Fd()))
}
<|endoftext|>