<commit_before>// Copyright (c) 2016 Pani Networks
// 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 listener listens to Kubernetes for policy updates.
package listener

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"

	"github.com/romana/core/common"
	"github.com/romana/core/tenant"
	log "github.com/romana/rlog"

	"k8s.io/client-go/1.5/kubernetes"
	"k8s.io/client-go/1.5/tools/cache"
	"k8s.io/client-go/1.5/tools/clientcmd"
)

type networkPolicyAction int

const (
	networkPolicyActionDelete networkPolicyAction = iota
	networkPolicyActionAdd
	networkPolicyActionModify
)

const (
	HttpGetParamWatch           = "watch=true"
	HttpGetParamResourceVersion = "resourceVersion"
)

// KubeListener is a Service that listens to updates
// from Kubernetes by connecting to the endpoints specified
// and consuming chunked JSON documents. The endpoints are
// constructed from kubeURL and the following paths:
// 1. namespaceNotificationPath for namespace additions/deletions
// 2. policyNotificationPathPrefix + <namespace name> + policyNotificationPathPostfix
//    for policy additions/deletions.
type KubeListener struct {
	config                        common.ServiceConfig
	restClient                    *common.RestClient
	kubeURL                       string
	namespaceNotificationPath     string
	policyNotificationPathPrefix  string
	policyNotificationPathPostfix string
	segmentLabelName              string
	tenantLabelName               string
	lastEventPerNamespace         map[string]uint64
	namespaceBufferSize           uint64

	kubeClient *kubernetes.Clientset
	Watchers   map[string]cache.ListerWatcher
}

// Routes returns various routes used in the service.
func (l *KubeListener) Routes() common.Routes {
	routes := common.Routes{}
	return routes
}

// Name implements method of Service interface.
func (l *KubeListener) Name() string {
	return "kubernetesListener"
}

func (l *KubeListener) CreateSchema(overwrite bool) error {
	return nil
}

// SetConfig implements SetConfig function of the Service interface.
func (l *KubeListener) SetConfig(config common.ServiceConfig) error {
	confString := "/etc/romana/romana.conf.yml:kubernetesListener:config:"
	log.Trace(5, confString, config)

	m := config.ServiceSpecific
	if kl, ok := m["kubernetes_url"]; !ok || kl == "" {
		return fmt.Errorf("%s%s", confString, "kubernetes_url required in config.")
	}
	l.kubeURL = m["kubernetes_url"].(string)

	if nnp, ok := m["namespace_notification_path"]; !ok || nnp == "" {
		return fmt.Errorf("%s%s", confString, "namespace_notification_path required in config.")
	}
	l.namespaceNotificationPath = m["namespace_notification_path"].(string)

	if pnppre, ok := m["policy_notification_path_prefix"]; !ok || pnppre == "" {
		return fmt.Errorf("%s%s", confString, "policy_notification_path_prefix required in config.")
	}
	l.policyNotificationPathPrefix = m["policy_notification_path_prefix"].(string)

	if pnppost, ok := m["policy_notification_path_prefix"]; !ok || pnppost == "" {
		return fmt.Errorf("%s%s", confString, "policy_notification_path_postfix required in config.")
	}
	l.policyNotificationPathPostfix = m["policy_notification_path_postfix"].(string)

	if sln, ok := m["segment_label_name"]; !ok || sln == "" {
		return fmt.Errorf("%s%s", confString, "segment_label_name required in config.")
	}
	l.segmentLabelName = m["segment_label_name"].(string)

	if tln, ok := m["tenant_label_name"]; !ok || tln == "" {
		return fmt.Errorf("%s%s", confString, "tenant_label_name required in config.")
	}
	l.tenantLabelName = m["tenant_label_name"].(string)

	l.namespaceBufferSize = 1000

	if kc, ok := m["kubernetes_config"]; !ok || kc == "" {
		// Default kubernetes config location on ubuntu
		// TODO: this should not be hard coded, other
		//       distributions may have other user names.
		m["kubernetes_config"] = "/home/ubuntu/.kube/config"
	}

	// TODO, this loads kubernetes config from flags provided in main
	// should be loading from path provided by romana-root. Stas.
	kubeClientConfig, err := clientcmd.BuildConfigFromFlags("", m["kubernetes_config"].(string))
	if err != nil {
		return errors.New(fmt.Sprintf("Failed to load kubernetes kubeClientConfig %s", err))
	}
	clientset, err := kubernetes.NewForConfig(kubeClientConfig)
	if err != nil {
		return fmt.Errorf("Failed to make kubernetes client %s", err)
	}
	l.kubeClient = clientset

	// TODO, find a better place to initialize
	// the translator. Stas.
	PTranslator.Init(l.restClient, l.segmentLabelName, l.tenantLabelName)
	tc := PTranslator.GetClient()
	if tc == nil {
		log.Critical("Failed to initialize rest client for policy translator.")
		os.Exit(255)
	}

	return nil
}

// TODO there should be a better way to introduce translator
// then global variable like this one.
var PTranslator Translator

// Run configures and runs listener service.
func Run(rootServiceURL string, cred *common.Credential) (*common.RestServiceInfo, error) {
	clientConfig := common.GetDefaultRestClientConfig(rootServiceURL)
	clientConfig.Credential = cred
	client, err := common.NewRestClient(clientConfig)
	if err != nil {
		return nil, err
	}
	KubeListener := &KubeListener{}
	KubeListener.restClient = client

	config, err := client.GetServiceConfig(KubeListener.Name())
	if err != nil {
		return nil, err
	}
	return common.InitializeService(KubeListener, *config, cred)
}

// getOrAddSegment finds a segment (based on segment selector).
// If not found, it adds one.
func (l *KubeListener) getOrAddSegment(namespace string, kubeSegmentName string) (*tenant.Segment, error) {
	ten := &tenant.Tenant{}
	ten.Name = namespace
	// TODO this should be changed to find EXACTLY one after deletion functionality is implemented
	err := l.restClient.Find(ten, common.FindLast)
	if err != nil {
		return nil, err
	}

	seg := &tenant.Segment{}
	seg.Name = kubeSegmentName
	seg.TenantID = ten.ID
	err = l.restClient.Find(seg, common.FindExactlyOne)
	if err == nil {
		return seg, nil
	}

	switch err := err.(type) {
	case common.HttpError:
		if err.StatusCode == http.StatusNotFound {
			// Not found, so let's create a segment.
			segreq := tenant.Segment{Name: kubeSegmentName, TenantID: ten.ID}
			segURL, err2 := l.restClient.GetServiceUrl("tenant")
			if err2 != nil {
				return nil, err2
			}
			segURL = fmt.Sprintf("%s/tenants/%d/segments", segURL, ten.ID)
			err2 = l.restClient.Post(segURL, segreq, seg)
			if err2 == nil {
				// Successful creation.
				return seg, nil
			}
			// Creation of non-existing segment gave an error.
			switch err2 := err2.(type) {
			case common.HttpError:
				// Maybe someone else just created a segment between the original
				// lookup and now?
				if err2.StatusCode == http.StatusConflict {
					switch details := err2.Details.(type) {
					case tenant.Segment:
						// We expect the existing segment to be returned in the details field.
						return &details, nil
					default:
						// This is unexpected...
						return nil, err
					}
				}
				// Any other HTTP error other than a Conflict here - return it.
				return nil, err2
			default:
				// Any other error - return it
				return nil, err2
			}
		}
		// Any other HTTP error other than a Not found here - return it
		return nil, err
	default:
		// Any other error - return it
		return nil, err
	}
}

// resolveTenantByName retrieves tenant information from romana.
func (l *KubeListener) resolveTenantByName(tenantName string) (*tenant.Tenant, error) {
	t := &tenant.Tenant{Name: tenantName}
	err := l.restClient.Find(t, common.FindLast)
	if err != nil {
		return t, err
	}
	return t, nil
}

func (l *KubeListener) applyNetworkPolicy(action networkPolicyAction, romanaNetworkPolicy common.Policy) error {
	policyURL, err := l.restClient.GetServiceUrl("policy")
	if err != nil {
		return err
	}
	policyURL = fmt.Sprintf("%s/policies", policyURL)
	policyStr, _ := json.Marshal(romanaNetworkPolicy)
	switch action {
	case networkPolicyActionAdd:
		log.Infof("Applying policy %s", policyStr)
		err := l.restClient.Post(policyURL, romanaNetworkPolicy, &romanaNetworkPolicy)
		if err != nil {
			return err
		}
	case networkPolicyActionDelete:
		log.Infof("Deleting policy policy %s", policyStr)
		err := l.restClient.Delete(policyURL, romanaNetworkPolicy, &romanaNetworkPolicy)
		if err != nil {
			return err
		}
	default:
		return errors.New("Unsupported operation")
	}
	return nil
}

func (l *KubeListener) Initialize(client *common.RestClient) error {
	l.restClient = client
	l.lastEventPerNamespace = make(map[string]uint64)
	log.Infof("%s: Starting server", l.Name())
	nsURL, err := common.CleanURL(fmt.Sprintf("%s/%s/?%s", l.kubeURL, l.namespaceNotificationPath, HttpGetParamWatch))
	if err != nil {
		return err
	}
	log.Infof("Starting to listen on %s", nsURL)
	done := make(chan struct{})
	eventc, err := l.nsWatch(done, nsURL)
	if err != nil {
		log.Critical("Namespace watcher failed to start", err)
		os.Exit(255)
	}

	// events := l.conductor(nsEvents, done)
	l.process(eventc, done)

	ProduceNewPolicyEvents(eventc, done, l)

	log.Info("All routines started")
	return nil
}

// CreateSchema is placeholder for now.
func CreateSchema(rootServiceURL string, overwrite bool) error {
	return nil
}
<commit_msg>listener: move translator initilization code in listener intialization.<commit_after>// Copyright (c) 2016 Pani Networks
// 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 listener listens to Kubernetes for policy updates.
package listener

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"

	"github.com/romana/core/common"
	"github.com/romana/core/tenant"
	log "github.com/romana/rlog"

	"k8s.io/client-go/1.5/kubernetes"
	"k8s.io/client-go/1.5/tools/cache"
	"k8s.io/client-go/1.5/tools/clientcmd"
)

type networkPolicyAction int

const (
	networkPolicyActionDelete networkPolicyAction = iota
	networkPolicyActionAdd
	networkPolicyActionModify
)

const (
	HttpGetParamWatch           = "watch=true"
	HttpGetParamResourceVersion = "resourceVersion"
)

// KubeListener is a Service that listens to updates
// from Kubernetes by connecting to the endpoints specified
// and consuming chunked JSON documents. The endpoints are
// constructed from kubeURL and the following paths:
// 1. namespaceNotificationPath for namespace additions/deletions
// 2. policyNotificationPathPrefix + <namespace name> + policyNotificationPathPostfix
//    for policy additions/deletions.
type KubeListener struct {
	config                        common.ServiceConfig
	restClient                    *common.RestClient
	kubeURL                       string
	namespaceNotificationPath     string
	policyNotificationPathPrefix  string
	policyNotificationPathPostfix string
	segmentLabelName              string
	tenantLabelName               string
	lastEventPerNamespace         map[string]uint64
	namespaceBufferSize           uint64

	kubeClient *kubernetes.Clientset
	Watchers   map[string]cache.ListerWatcher
}

// Routes returns various routes used in the service.
func (l *KubeListener) Routes() common.Routes {
	routes := common.Routes{}
	return routes
}

// Name implements method of Service interface.
func (l *KubeListener) Name() string {
	return "kubernetesListener"
}

func (l *KubeListener) CreateSchema(overwrite bool) error {
	return nil
}

// SetConfig implements SetConfig function of the Service interface.
func (l *KubeListener) SetConfig(config common.ServiceConfig) error {
	confString := "/etc/romana/romana.conf.yml:kubernetesListener:config:"
	log.Trace(5, confString, config)

	m := config.ServiceSpecific
	if kl, ok := m["kubernetes_url"]; !ok || kl == "" {
		return fmt.Errorf("%s%s", confString, "kubernetes_url required in config.")
	}
	l.kubeURL = m["kubernetes_url"].(string)

	if nnp, ok := m["namespace_notification_path"]; !ok || nnp == "" {
		return fmt.Errorf("%s%s", confString, "namespace_notification_path required in config.")
	}
	l.namespaceNotificationPath = m["namespace_notification_path"].(string)

	if pnppre, ok := m["policy_notification_path_prefix"]; !ok || pnppre == "" {
		return fmt.Errorf("%s%s", confString, "policy_notification_path_prefix required in config.")
	}
	l.policyNotificationPathPrefix = m["policy_notification_path_prefix"].(string)

	if pnppost, ok := m["policy_notification_path_prefix"]; !ok || pnppost == "" {
		return fmt.Errorf("%s%s", confString, "policy_notification_path_postfix required in config.")
	}
	l.policyNotificationPathPostfix = m["policy_notification_path_postfix"].(string)

	if sln, ok := m["segment_label_name"]; !ok || sln == "" {
		return fmt.Errorf("%s%s", confString, "segment_label_name required in config.")
	}
	l.segmentLabelName = m["segment_label_name"].(string)

	if tln, ok := m["tenant_label_name"]; !ok || tln == "" {
		return fmt.Errorf("%s%s", confString, "tenant_label_name required in config.")
	}
	l.tenantLabelName = m["tenant_label_name"].(string)

	l.namespaceBufferSize = 1000

	if kc, ok := m["kubernetes_config"]; !ok || kc == "" {
		// Default kubernetes config location on ubuntu
		// TODO: this should not be hard coded, other
		//       distributions may have other user names.
		m["kubernetes_config"] = "/home/ubuntu/.kube/config"
	}

	// TODO, this loads kubernetes config from flags provided in main
	// should be loading from path provided by romana-root. Stas.
	kubeClientConfig, err := clientcmd.BuildConfigFromFlags("", m["kubernetes_config"].(string))
	if err != nil {
		return errors.New(fmt.Sprintf("Failed to load kubernetes kubeClientConfig %s", err))
	}
	clientset, err := kubernetes.NewForConfig(kubeClientConfig)
	if err != nil {
		return fmt.Errorf("Failed to make kubernetes client %s", err)
	}
	l.kubeClient = clientset

	return nil
}

// TODO there should be a better way to introduce translator
// then global variable like this one.
var PTranslator Translator

// Run configures and runs listener service.
func Run(rootServiceURL string, cred *common.Credential) (*common.RestServiceInfo, error) {
	clientConfig := common.GetDefaultRestClientConfig(rootServiceURL)
	clientConfig.Credential = cred
	client, err := common.NewRestClient(clientConfig)
	if err != nil {
		return nil, err
	}
	KubeListener := &KubeListener{}
	KubeListener.restClient = client

	config, err := client.GetServiceConfig(KubeListener.Name())
	if err != nil {
		return nil, err
	}
	return common.InitializeService(KubeListener, *config, cred)
}

// getOrAddSegment finds a segment (based on segment selector).
// If not found, it adds one.
func (l *KubeListener) getOrAddSegment(namespace string, kubeSegmentName string) (*tenant.Segment, error) {
	ten := &tenant.Tenant{}
	ten.Name = namespace
	// TODO this should be changed to find EXACTLY one after deletion functionality is implemented
	err := l.restClient.Find(ten, common.FindLast)
	if err != nil {
		return nil, err
	}

	seg := &tenant.Segment{}
	seg.Name = kubeSegmentName
	seg.TenantID = ten.ID
	err = l.restClient.Find(seg, common.FindExactlyOne)
	if err == nil {
		return seg, nil
	}

	switch err := err.(type) {
	case common.HttpError:
		if err.StatusCode == http.StatusNotFound {
			// Not found, so let's create a segment.
			segreq := tenant.Segment{Name: kubeSegmentName, TenantID: ten.ID}
			segURL, err2 := l.restClient.GetServiceUrl("tenant")
			if err2 != nil {
				return nil, err2
			}
			segURL = fmt.Sprintf("%s/tenants/%d/segments", segURL, ten.ID)
			err2 = l.restClient.Post(segURL, segreq, seg)
			if err2 == nil {
				// Successful creation.
				return seg, nil
			}
			// Creation of non-existing segment gave an error.
			switch err2 := err2.(type) {
			case common.HttpError:
				// Maybe someone else just created a segment between the original
				// lookup and now?
				if err2.StatusCode == http.StatusConflict {
					switch details := err2.Details.(type) {
					case tenant.Segment:
						// We expect the existing segment to be returned in the details field.
						return &details, nil
					default:
						// This is unexpected...
						return nil, err
					}
				}
				// Any other HTTP error other than a Conflict here - return it.
				return nil, err2
			default:
				// Any other error - return it
				return nil, err2
			}
		}
		// Any other HTTP error other than a Not found here - return it
		return nil, err
	default:
		// Any other error - return it
		return nil, err
	}
}

// resolveTenantByName retrieves tenant information from romana.
func (l *KubeListener) resolveTenantByName(tenantName string) (*tenant.Tenant, error) {
	t := &tenant.Tenant{Name: tenantName}
	err := l.restClient.Find(t, common.FindLast)
	if err != nil {
		return t, err
	}
	return t, nil
}

func (l *KubeListener) applyNetworkPolicy(action networkPolicyAction, romanaNetworkPolicy common.Policy) error {
	policyURL, err := l.restClient.GetServiceUrl("policy")
	if err != nil {
		return err
	}
	policyURL = fmt.Sprintf("%s/policies", policyURL)
	policyStr, _ := json.Marshal(romanaNetworkPolicy)
	switch action {
	case networkPolicyActionAdd:
		log.Infof("Applying policy %s", policyStr)
		err := l.restClient.Post(policyURL, romanaNetworkPolicy, &romanaNetworkPolicy)
		if err != nil {
			return err
		}
	case networkPolicyActionDelete:
		log.Infof("Deleting policy policy %s", policyStr)
		err := l.restClient.Delete(policyURL, romanaNetworkPolicy, &romanaNetworkPolicy)
		if err != nil {
			return err
		}
	default:
		return errors.New("Unsupported operation")
	}
	return nil
}

func (l *KubeListener) Initialize(client *common.RestClient) error {
	l.restClient = client

	// TODO, find a better place to initialize
	// the translator. Stas.
	PTranslator.Init(l.restClient, l.segmentLabelName, l.tenantLabelName)
	tc := PTranslator.GetClient()
	if tc == nil {
		log.Critical("Failed to initialize rest client for policy translator.")
		os.Exit(255)
	}

	l.lastEventPerNamespace = make(map[string]uint64)
	log.Infof("%s: Starting server", l.Name())
	nsURL, err := common.CleanURL(fmt.Sprintf("%s/%s/?%s", l.kubeURL, l.namespaceNotificationPath, HttpGetParamWatch))
	if err != nil {
		return err
	}
	log.Infof("Starting to listen on %s", nsURL)
	done := make(chan struct{})
	eventc, err := l.nsWatch(done, nsURL)
	if err != nil {
		log.Critical("Namespace watcher failed to start", err)
		os.Exit(255)
	}

	// events := l.conductor(nsEvents, done)
	l.process(eventc, done)

	ProduceNewPolicyEvents(eventc, done, l)

	log.Info("All routines started")
	return nil
}

// CreateSchema is placeholder for now.
func CreateSchema(rootServiceURL string, overwrite bool) error {
	return nil
}
<|endoftext|>


<commit_before>package http

import (
	"fmt"
	"github.com/antonholmquist/jason"
	"k8-rolling-demo/lib/model"
	"net/http"
	"os"
	"sync"
)

type HttpClient struct {
	host      string
	port      string
	namespace string
}

func NewHttpClient(host, port, namespace string) *HttpClient {
	return &HttpClient{
		host:      host,
		port:      port,
		namespace: namespace,
	}
}

func DefaultHttpClient() *HttpClient {
	host := os.Getenv("K8_API_HOST")
	port := os.Getenv("K8_API_PORT")
	namespace := "k8-rolling-demo"

	return NewHttpClient(host, port, namespace)
}

func (c *HttpClient) FetchServiceDetails(service string) (*model.Service, error) {
	uri := fmt.Sprintf(
		"http://%s:%s/api/v1/namespaces/%s/services/%s",
		c.host,
		c.port,
		c.namespace,
		service,
	)

	resp, err := http.Get(uri)
	if err != nil {
		return nil, err
	}

	v, err := jason.NewObjectFromReader(resp.Body)
	if err != nil {
		return nil, err
	}

	uid, _ := v.GetString("metadata", "uid")
	name, _ := v.GetString("metadata", "name")

	labels := make(map[string]string)
	labelsObject, _ := v.GetObject("metadata", "labels")
	for key, value := range labelsObject.Map() {
		str, _ := value.String()
		labels[key] = str
	}

	selectors := make(map[string]string)
	selectorsObject, _ := v.GetObject("spec", "selector")
	for key, value := range selectorsObject.Map() {
		str, _ := value.String()
		selectors[key] = str
	}

	svc := model.NewService(uid, name, labels, selectors)

	podNames, err := c.fetchPodNamesForService(name)
	if err != nil {
		return nil, err
	}

	wg := &sync.WaitGroup{}
	wg.Add(len(podNames))

	for _, podName := range podNames {
		go func(podName string, wg *sync.WaitGroup) {
			pod, _ := c.FetchPodDetails(podName)
			svc.AddPod(pod)

			wg.Done()
		}(podName, wg)
	}

	wg.Wait()

	return svc, nil
}

func (c *HttpClient) fetchPodNamesForService(service string) ([]string, error) {
	uri := fmt.Sprintf(
		"http://%s:%s/api/v1/namespaces/%s/endpoints/%s",
		c.host,
		c.port,
		c.namespace,
		service,
	)

	resp, err := http.Get(uri)
	if err != nil {
		return nil, err
	}

	v, err := jason.NewObjectFromReader(resp.Body)
	if err != nil {
		return nil, err
	}

	subsets, _ := v.GetObjectArray("subsets")
	addrs, _ := subsets[0].GetObjectArray("addresses")

	podNames := make([]string, 0)
	for _, addr := range addrs {
		podName, _ := addr.GetString("targetRef", "name")
		podNames = append(podNames, podName)
	}

	return podNames, nil
}

func (c *HttpClient) FetchPodDetails(pod string) (*model.Pod, error) {
	uri := fmt.Sprintf(
		"http://%s:%s/api/v1/namespaces/%s/pods/%s",
		c.host,
		c.port,
		c.namespace,
		pod,
	)

	resp, err := http.Get(uri)
	if err != nil {
		return nil, err
	}

	v, err := jason.NewObjectFromReader(resp.Body)
	if err != nil {
		return nil, err
	}

	uid, _ := v.GetString("metadata", "uid")
	name, _ := v.GetString("metadata", "name")
	ipAddr, _ := v.GetString("status", "podIp")

	labels := make(map[string]string)
	labelsObject, _ := v.GetObject("metadata", "labels")
	for key, value := range labelsObject.Map() {
		str, _ := value.String()
		labels[key] = str
	}

	status, _ := v.GetString("status", "phase")
	conditions, _ := v.GetObjectArray("status", "conditions")
	firstCondition, _ := conditions[0].GetString("type")

	pd := &model.Pod{
		Uid:       uid,
		Name:      name,
		IpAddr:    ipAddr,
		Labels:    labels,
		Status:    status,
		Condition: firstCondition,
	}

	return pd, nil
}
<commit_msg>Bugfix: Catches some conditions which caused panics<commit_after>package http

import (
	"errors"
	"fmt"
	"github.com/antonholmquist/jason"
	"k8-rolling-demo/lib/model"
	"net/http"
	"os"
	"sync"
)

type HttpClient struct {
	host      string
	port      string
	namespace string
}

func NewHttpClient(host, port, namespace string) *HttpClient {
	return &HttpClient{
		host:      host,
		port:      port,
		namespace: namespace,
	}
}

func DefaultHttpClient() *HttpClient {
	host := os.Getenv("K8_API_HOST")
	port := os.Getenv("K8_API_PORT")
	namespace := "k8-rolling-demo"

	return NewHttpClient(host, port, namespace)
}

func (c *HttpClient) FetchServiceDetails(service string) (*model.Service, error) {
	uri := fmt.Sprintf(
		"http://%s:%s/api/v1/namespaces/%s/services/%s",
		c.host,
		c.port,
		c.namespace,
		service,
	)

	resp, err := http.Get(uri)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode >= 400 {
		return nil, errors.New(resp.Status)
	}

	v, err := jason.NewObjectFromReader(resp.Body)
	if err != nil {
		return nil, err
	}

	uid, _ := v.GetString("metadata", "uid")
	name, _ := v.GetString("metadata", "name")

	labels := make(map[string]string)
	labelsObject, _ := v.GetObject("metadata", "labels")
	for key, value := range labelsObject.Map() {
		str, _ := value.String()
		labels[key] = str
	}

	selectors := make(map[string]string)
	selectorsObject, _ := v.GetObject("spec", "selector")
	for key, value := range selectorsObject.Map() {
		str, _ := value.String()
		selectors[key] = str
	}

	svc := model.NewService(uid, name, labels, selectors)

	podNames, err := c.fetchPodNamesForService(name)
	if err != nil {
		return nil, err
	}

	wg := &sync.WaitGroup{}
	wg.Add(len(podNames))

	for _, podName := range podNames {
		go func(podName string, wg *sync.WaitGroup) {
			pod, _ := c.FetchPodDetails(podName)
			svc.AddPod(pod)

			wg.Done()
		}(podName, wg)
	}

	wg.Wait()

	return svc, nil
}

func (c *HttpClient) fetchPodNamesForService(service string) ([]string, error) {
	uri := fmt.Sprintf(
		"http://%s:%s/api/v1/namespaces/%s/endpoints/%s",
		c.host,
		c.port,
		c.namespace,
		service,
	)

	resp, err := http.Get(uri)
	if err != nil {
		return nil, err
	}

	v, err := jason.NewObjectFromReader(resp.Body)
	if err != nil {
		return nil, err
	}

	subsets, _ := v.GetObjectArray("subsets")
	if len(subsets) <= 0 {
		return nil, errors.New("No pods available in service")
	}

	addrs, _ := subsets[0].GetObjectArray("addresses")

	podNames := make([]string, 0)
	for _, addr := range addrs {
		podName, _ := addr.GetString("targetRef", "name")
		podNames = append(podNames, podName)
	}

	return podNames, nil
}

func (c *HttpClient) FetchPodDetails(pod string) (*model.Pod, error) {
	uri := fmt.Sprintf(
		"http://%s:%s/api/v1/namespaces/%s/pods/%s",
		c.host,
		c.port,
		c.namespace,
		pod,
	)

	resp, err := http.Get(uri)
	if err != nil {
		return nil, err
	}

	v, err := jason.NewObjectFromReader(resp.Body)
	if err != nil {
		return nil, err
	}

	uid, _ := v.GetString("metadata", "uid")
	name, _ := v.GetString("metadata", "name")
	ipAddr, _ := v.GetString("status", "podIp")

	labels := make(map[string]string)
	labelsObject, _ := v.GetObject("metadata", "labels")
	for key, value := range labelsObject.Map() {
		str, _ := value.String()
		labels[key] = str
	}

	status, _ := v.GetString("status", "phase")
	conditions, _ := v.GetObjectArray("status", "conditions")
	firstCondition, _ := conditions[0].GetString("type")

	pd := &model.Pod{
		Uid:       uid,
		Name:      name,
		IpAddr:    ipAddr,
		Labels:    labels,
		Status:    status,
		Condition: firstCondition,
	}

	return pd, nil
}
<|endoftext|>


<commit_before>package mackerel

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"
)

// Host host information
type Host struct {
	ID               string      `json:"id"`
	Name             string      `json:"name"`
	DisplayName      string      `json:"displayName,omitempty"`
	CustomIdentifier string      `json:"customIdentifier,omitempty"`
	Type             string      `json:"type"`
	Status           string      `json:"status"`
	Memo             string      `json:"memo"`
	Roles            Roles       `json:"roles"`
	IsRetired        bool        `json:"isRetired"`
	CreatedAt        int32       `json:"createdAt"`
	Meta             HostMeta    `json:"meta"`
	Interfaces       []Interface `json:"interfaces"`
}

// Roles host role maps
type Roles map[string][]string

// HostMeta host meta informations
type HostMeta struct {
	AgentRevision string      `json:"agent-revision,omitempty"`
	AgentVersion  string      `json:"agent-version,omitempty"`
	AgentName     string      `json:"agent-name,omitempty"`
	BlockDevice   BlockDevice `json:"block_device,omitempty"`
	CPU           CPU         `json:"cpu,omitempty"`
	Filesystem    FileSystem  `json:"filesystem,omitempty"`
	Kernel        Kernel      `json:"kernel,omitempty"`
	Memory        Memory      `json:"memory,omitempty"`
	Cloud         *Cloud      `json:"cloud,omitempty"`
}

// BlockDevice blockdevice
type BlockDevice map[string]map[string]interface{}

// CPU cpu
type CPU []map[string]interface{}

// FileSystem filesystem
type FileSystem map[string]interface{}

// Kernel kernel
type Kernel map[string]string

// Memory memory
type Memory map[string]string

// Cloud cloud
type Cloud struct {
	Provider string      `json:"provider,omitempty"`
	MetaData interface{} `json:"metadata,omitempty"`
}

// Interface network interface
type Interface struct {
	Name          string   `json:"name,omitempty"`
	IPAddress     string   `json:"ipAddress,omitempty"`
	IPv4Addresses []string `json:"ipv4Addresses,omitempty"`
	IPv6Addresses []string `json:"ipv6Addresses,omitempty"`
	MacAddress    string   `json:"macAddress,omitempty"`
}

// FindHostsParam parameters for FindHosts
type FindHostsParam struct {
	Service          string
	Roles            []string
	Name             string
	Statuses         []string
	CustomIdentifier string
}

// CreateHostParam parameters for CreateHost
type CreateHostParam struct {
	Name             string        `json:"name"`
	DisplayName      string        `json:"displayName,omitempty"`
	Meta             HostMeta      `json:"meta"`
	Interfaces       []Interface   `json:"interfaces"`
	RoleFullnames    []string      `json:"roleFullnames"`
	Checks           []CheckConfig `json:"checks"`
	CustomIdentifier string        `json:"customIdentifier,omitempty"`
}

// CheckConfig is check plugin name and memo
type CheckConfig struct {
	Name string `json:"name,omitempty"`
	Memo string `json:"memo,omitempty"`
}

// UpdateHostParam parameters for UpdateHost
type UpdateHostParam CreateHostParam

const (
	// HostStatusWorking represents "working" status
	HostStatusWorking = "working"
	// HostStatusStandby represents "standby" status
	HostStatusStandby = "standby"
	// HostStatusMaintenance represents "maintenance" status
	HostStatusMaintenance = "maintenance"
	// HostStatusPoweroff represents "poeroff" status
	HostStatusPoweroff = "poweroff"
)

// GetRoleFullnames getrolefullnames
func (h *Host) GetRoleFullnames() []string {
	if len(h.Roles) < 1 {
		return nil
	}

	var fullnames []string
	for service, roles := range h.Roles {
		for _, role := range roles {
			fullname := strings.Join([]string{service, role}, ":")
			fullnames = append(fullnames, fullname)
		}
	}

	return fullnames
}

// DateFromCreatedAt returns time.Time
func (h *Host) DateFromCreatedAt() time.Time {
	return time.Unix(int64(h.CreatedAt), 0)
}

// IPAddresses returns ipaddresses
func (h *Host) IPAddresses() map[string]string {
	if len(h.Interfaces) < 1 {
		return nil
	}

	ipAddresses := make(map[string]string, 0)
	for _, iface := range h.Interfaces {
		ipAddresses[iface.Name] = iface.IPAddress
	}
	return ipAddresses
}

// FindHost find the host
func (c *Client) FindHost(id string) (*Host, error) {
	req, err := http.NewRequest("GET", c.urlFor(fmt.Sprintf("/api/v0/hosts/%s", id)).String(), nil)
	if err != nil {
		return nil, err
	}
	resp, err := c.Request(req)
	defer closeResponse(resp)
	if err != nil {
		return nil, err
	}

	var data struct {
		Host *Host `json:"host"`
	}
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return nil, err
	}
	return data.Host, err
}

// FindHosts find hosts
func (c *Client) FindHosts(param *FindHostsParam) ([]*Host, error) {
	v := url.Values{}
	if param.Service != "" {
		v.Set("service", param.Service)
	}
	if len(param.Roles) >= 1 {
		for _, role := range param.Roles {
			v.Add("role", role)
		}
	}
	if param.Name != "" {
		v.Set("name", param.Name)
	}
	if len(param.Statuses) >= 1 {
		for _, status := range param.Statuses {
			v.Add("status", status)
		}
	}
	if param.CustomIdentifier != "" {
		v.Set("customIdentifier", param.CustomIdentifier)
	}

	req, err := http.NewRequest("GET", fmt.Sprintf("%s?%s", c.urlFor("/api/v0/hosts").String(), v.Encode()), nil)
	if err != nil {
		return nil, err
	}
	resp, err := c.Request(req)
	defer closeResponse(resp)
	if err != nil {
		return nil, err
	}

	var data struct {
		Hosts []*Host `json:"hosts"`
	}
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return nil, err
	}

	return data.Hosts, err
}

// CreateHost creating host
func (c *Client) CreateHost(param *CreateHostParam) (string, error) {
	resp, err := c.PostJSON("/api/v0/hosts", param)
	defer closeResponse(resp)
	if err != nil {
		return "", err
	}

	var data struct {
		ID string `json:"id"`
	}
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return "", err
	}
	return data.ID, nil
}

// UpdateHost updates host
func (c *Client) UpdateHost(hostID string, param *UpdateHostParam) (string, error) {
	resp, err := c.PutJSON(fmt.Sprintf("/api/v0/hosts/%s", hostID), param)
	defer closeResponse(resp)
	if err != nil {
		return "", err
	}

	var data struct {
		ID string `json:"id"`
	}
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return "", err
	}

	return data.ID, nil
}

// UpdateHostStatus updates host status
func (c *Client) UpdateHostStatus(hostID string, status string) error {
	resp, err := c.PostJSON(fmt.Sprintf("/api/v0/hosts/%s/status", hostID), map[string]string{
		"status": status,
	})
	defer closeResponse(resp)
	if err != nil {
		return err
	}
	return nil
}

// UpdateHostRoleFullnames updates host roles
func (c *Client) UpdateHostRoleFullnames(hostID string, roleFullnames []string) error {
	resp, err := c.PutJSON(fmt.Sprintf("/api/v0/hosts/%s/role-fullnames", hostID), map[string][]string{
		"roleFullnames": roleFullnames,
	})
	defer closeResponse(resp)
	if err != nil {
		return err
	}
	return nil
}

// RetireHost retuire the host
func (c *Client) RetireHost(id string) error {
	resp, err := c.PostJSON(fmt.Sprintf("/api/v0/hosts/%s/retire", id), "{}")
	defer closeResponse(resp)
	return err
}

// ListHostMetricNames lists metric names of a host
func (c *Client) ListHostMetricNames(id string) ([]string, error) {
	req, err := http.NewRequest("GET", c.urlFor(fmt.Sprintf("/api/v0/hosts/%s/metric-names", id)).String(), nil)
	if err != nil {
		return nil, err
	}
	resp, err := c.Request(req)
	defer closeResponse(resp)
	if err != nil {
		return nil, err
	}

	var data struct {
		Names []string `json:"names"`
	}
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return nil, err
	}
	return data.Names, err
}
<commit_msg>Fix typo.<commit_after>package mackerel

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"
)

// Host host information
type Host struct {
	ID               string      `json:"id"`
	Name             string      `json:"name"`
	DisplayName      string      `json:"displayName,omitempty"`
	CustomIdentifier string      `json:"customIdentifier,omitempty"`
	Type             string      `json:"type"`
	Status           string      `json:"status"`
	Memo             string      `json:"memo"`
	Roles            Roles       `json:"roles"`
	IsRetired        bool        `json:"isRetired"`
	CreatedAt        int32       `json:"createdAt"`
	Meta             HostMeta    `json:"meta"`
	Interfaces       []Interface `json:"interfaces"`
}

// Roles host role maps
type Roles map[string][]string

// HostMeta host meta informations
type HostMeta struct {
	AgentRevision string      `json:"agent-revision,omitempty"`
	AgentVersion  string      `json:"agent-version,omitempty"`
	AgentName     string      `json:"agent-name,omitempty"`
	BlockDevice   BlockDevice `json:"block_device,omitempty"`
	CPU           CPU         `json:"cpu,omitempty"`
	Filesystem    FileSystem  `json:"filesystem,omitempty"`
	Kernel        Kernel      `json:"kernel,omitempty"`
	Memory        Memory      `json:"memory,omitempty"`
	Cloud         *Cloud      `json:"cloud,omitempty"`
}

// BlockDevice blockdevice
type BlockDevice map[string]map[string]interface{}

// CPU cpu
type CPU []map[string]interface{}

// FileSystem filesystem
type FileSystem map[string]interface{}

// Kernel kernel
type Kernel map[string]string

// Memory memory
type Memory map[string]string

// Cloud cloud
type Cloud struct {
	Provider string      `json:"provider,omitempty"`
	MetaData interface{} `json:"metadata,omitempty"`
}

// Interface network interface
type Interface struct {
	Name          string   `json:"name,omitempty"`
	IPAddress     string   `json:"ipAddress,omitempty"`
	IPv4Addresses []string `json:"ipv4Addresses,omitempty"`
	IPv6Addresses []string `json:"ipv6Addresses,omitempty"`
	MacAddress    string   `json:"macAddress,omitempty"`
}

// FindHostsParam parameters for FindHosts
type FindHostsParam struct {
	Service          string
	Roles            []string
	Name             string
	Statuses         []string
	CustomIdentifier string
}

// CreateHostParam parameters for CreateHost
type CreateHostParam struct {
	Name             string        `json:"name"`
	DisplayName      string        `json:"displayName,omitempty"`
	Meta             HostMeta      `json:"meta"`
	Interfaces       []Interface   `json:"interfaces"`
	RoleFullnames    []string      `json:"roleFullnames"`
	Checks           []CheckConfig `json:"checks"`
	CustomIdentifier string        `json:"customIdentifier,omitempty"`
}

// CheckConfig is check plugin name and memo
type CheckConfig struct {
	Name string `json:"name,omitempty"`
	Memo string `json:"memo,omitempty"`
}

// UpdateHostParam parameters for UpdateHost
type UpdateHostParam CreateHostParam

const (
	// HostStatusWorking represents "working" status
	HostStatusWorking = "working"
	// HostStatusStandby represents "standby" status
	HostStatusStandby = "standby"
	// HostStatusMaintenance represents "maintenance" status
	HostStatusMaintenance = "maintenance"
	// HostStatusPoweroff represents "poeroff" status
	HostStatusPoweroff = "poweroff"
)

// GetRoleFullnames gets role-full-names
func (h *Host) GetRoleFullnames() []string {
	if len(h.Roles) < 1 {
		return nil
	}

	var fullnames []string
	for service, roles := range h.Roles {
		for _, role := range roles {
			fullname := strings.Join([]string{service, role}, ":")
			fullnames = append(fullnames, fullname)
		}
	}

	return fullnames
}

// DateFromCreatedAt returns time.Time
func (h *Host) DateFromCreatedAt() time.Time {
	return time.Unix(int64(h.CreatedAt), 0)
}

// IPAddresses returns ipaddresses
func (h *Host) IPAddresses() map[string]string {
	if len(h.Interfaces) < 1 {
		return nil
	}

	ipAddresses := make(map[string]string, 0)
	for _, iface := range h.Interfaces {
		ipAddresses[iface.Name] = iface.IPAddress
	}
	return ipAddresses
}

// FindHost finds the host
func (c *Client) FindHost(id string) (*Host, error) {
	req, err := http.NewRequest("GET", c.urlFor(fmt.Sprintf("/api/v0/hosts/%s", id)).String(), nil)
	if err != nil {
		return nil, err
	}
	resp, err := c.Request(req)
	defer closeResponse(resp)
	if err != nil {
		return nil, err
	}

	var data struct {
		Host *Host `json:"host"`
	}
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return nil, err
	}
	return data.Host, err
}

// FindHosts finds hosts
func (c *Client) FindHosts(param *FindHostsParam) ([]*Host, error) {
	v := url.Values{}
	if param.Service != "" {
		v.Set("service", param.Service)
	}
	if len(param.Roles) >= 1 {
		for _, role := range param.Roles {
			v.Add("role", role)
		}
	}
	if param.Name != "" {
		v.Set("name", param.Name)
	}
	if len(param.Statuses) >= 1 {
		for _, status := range param.Statuses {
			v.Add("status", status)
		}
	}
	if param.CustomIdentifier != "" {
		v.Set("customIdentifier", param.CustomIdentifier)
	}

	req, err := http.NewRequest("GET", fmt.Sprintf("%s?%s", c.urlFor("/api/v0/hosts").String(), v.Encode()), nil)
	if err != nil {
		return nil, err
	}
	resp, err := c.Request(req)
	defer closeResponse(resp)
	if err != nil {
		return nil, err
	}

	var data struct {
		Hosts []*Host `json:"hosts"`
	}
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return nil, err
	}

	return data.Hosts, err
}

// CreateHost creates host
func (c *Client) CreateHost(param *CreateHostParam) (string, error) {
	resp, err := c.PostJSON("/api/v0/hosts", param)
	defer closeResponse(resp)
	if err != nil {
		return "", err
	}

	var data struct {
		ID string `json:"id"`
	}
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return "", err
	}
	return data.ID, nil
}

// UpdateHost updates host
func (c *Client) UpdateHost(hostID string, param *UpdateHostParam) (string, error) {
	resp, err := c.PutJSON(fmt.Sprintf("/api/v0/hosts/%s", hostID), param)
	defer closeResponse(resp)
	if err != nil {
		return "", err
	}

	var data struct {
		ID string `json:"id"`
	}
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return "", err
	}

	return data.ID, nil
}

// UpdateHostStatus updates host status
func (c *Client) UpdateHostStatus(hostID string, status string) error {
	resp, err := c.PostJSON(fmt.Sprintf("/api/v0/hosts/%s/status", hostID), map[string]string{
		"status": status,
	})
	defer closeResponse(resp)
	if err != nil {
		return err
	}
	return nil
}

// UpdateHostRoleFullnames updates host roles
func (c *Client) UpdateHostRoleFullnames(hostID string, roleFullnames []string) error {
	resp, err := c.PutJSON(fmt.Sprintf("/api/v0/hosts/%s/role-fullnames", hostID), map[string][]string{
		"roleFullnames": roleFullnames,
	})
	defer closeResponse(resp)
	if err != nil {
		return err
	}
	return nil
}

// RetireHost retires the host
func (c *Client) RetireHost(id string) error {
	resp, err := c.PostJSON(fmt.Sprintf("/api/v0/hosts/%s/retire", id), "{}")
	defer closeResponse(resp)
	return err
}

// ListHostMetricNames lists metric names of a host
func (c *Client) ListHostMetricNames(id string) ([]string, error) {
	req, err := http.NewRequest("GET", c.urlFor(fmt.Sprintf("/api/v0/hosts/%s/metric-names", id)).String(), nil)
	if err != nil {
		return nil, err
	}
	resp, err := c.Request(req)
	defer closeResponse(resp)
	if err != nil {
		return nil, err
	}

	var data struct {
		Names []string `json:"names"`
	}
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return nil, err
	}
	return data.Names, err
}
<|endoftext|>


<commit_before>package fm

import (
	"fmt"
	"go/ast"
	"go/token"
)

const (
	recvName = "f"
)

// FuncImplementer accepts an interface and returns implementations
// of its functions
type FuncImplementer interface {
	Implement(name *ast.Ident, i *ast.InterfaceType) []*ast.FuncDecl
}

// SpyFuncImplementer creates spy implementations of an interface's functions.
// Meant to be used in conjunction with SpyStructConverter
type SpyFuncImplementer struct{}

// Implement returns a function declaration whose arguments are saved
// as properties and whose return values are properties on a spy struct
func (s *SpyFuncImplementer) Implement(name *ast.Ident, i *ast.InterfaceType) []*ast.FuncDecl {
	var funcDecls []*ast.FuncDecl
	for _, list := range i.Methods.List {
		recv := &ast.FieldList{
			List: []*ast.Field{
				{
					Names: []*ast.Ident{ast.NewIdent(recvName)},
					Type:  &ast.StarExpr{X: name},
				},
			},
		}

		funcType, ok := list.Type.(*ast.FuncType)
		if !ok {
			// TODO: When will this happen?
			continue
		}

		funcDecls = append(funcDecls, &ast.FuncDecl{
			Recv: recv,
			Name: list.Names[0],
			Type: funcType,
			Body: createBlockStmt(list.Names[0].Name, funcType),
		})
	}
	return funcDecls
}

func createBlockStmt(fname string, f *ast.FuncType) *ast.BlockStmt {
	var list []ast.Stmt

	// add called assignment statement
	calledStmt := &ast.AssignStmt{
		Lhs: []ast.Expr{
			&ast.SelectorExpr{
				X:   ast.NewIdent(recvName),
				Sel: ast.NewIdent(fname + "_Called"),
			},
		},
		Tok: token.ASSIGN,
		Rhs: []ast.Expr{
			ast.NewIdent("true"),
		},
	}
	list = append(list, calledStmt)

	// add assignment for each param
	offset := 0
	for idx, field := range f.Params.List {
		// make assignments for multiple fields with same type
		if len(field.Names) > 1 {
			for _, name := range field.Names {
				assignStmt := &ast.AssignStmt{
					Lhs: []ast.Expr{
						&ast.SelectorExpr{
							X: &ast.SelectorExpr{
								X:   ast.NewIdent(recvName),
								Sel: ast.NewIdent(fname + inputSuffix),
							},
							Sel: ast.NewIdent(fmt.Sprintf("%s%d", argPrefix, idx+offset)),
						},
					},
					Tok: token.ASSIGN,
					Rhs: []ast.Expr{name},
				}
				list = append(list, assignStmt)

				offset++
			}
		} else {
			assignStmt := &ast.AssignStmt{
				Lhs: []ast.Expr{
					&ast.SelectorExpr{
						X: &ast.SelectorExpr{
							X:   ast.NewIdent(recvName),
							Sel: ast.NewIdent(fname + inputSuffix),
						},
						Sel: ast.NewIdent(fmt.Sprintf("%s%d", argPrefix, idx+offset)),
					},
				},
				Tok: token.ASSIGN,
				Rhs: []ast.Expr{field.Names[0]},
			}
			list = append(list, assignStmt)
		}
	}

	// add return statement if there are values to return
	var results []ast.Expr
	for idx := range f.Results.List {
		results = append(results, &ast.SelectorExpr{
			X: &ast.SelectorExpr{
				X:   ast.NewIdent(recvName), // for spy
				Sel: ast.NewIdent(fname + outputSuffix),
			},
			Sel: ast.NewIdent(fmt.Sprintf("%s%d", retPrefix, idx)),
		})
	}
	if len(results) > 0 {
		list = append(list, &ast.ReturnStmt{Results: results})
	}

	return &ast.BlockStmt{List: list}
}
<commit_msg>Formatting<commit_after>package fm

import (
	"fmt"
	"go/ast"
	"go/token"
)

const (
	recvName = "f"
)

// FuncImplementer accepts an interface and returns implementations
// of its functions
type FuncImplementer interface {
	Implement(name *ast.Ident, i *ast.InterfaceType) []*ast.FuncDecl
}

// SpyFuncImplementer creates spy implementations of an interface's functions.
// Meant to be used in conjunction with SpyStructConverter
type SpyFuncImplementer struct{}

// Implement returns a function declaration whose arguments are saved
// as properties and whose return values are properties on a spy struct
func (s *SpyFuncImplementer) Implement(name *ast.Ident, i *ast.InterfaceType) []*ast.FuncDecl {
	var funcDecls []*ast.FuncDecl
	for _, list := range i.Methods.List {
		recv := &ast.FieldList{
			List: []*ast.Field{{
				Names: []*ast.Ident{ast.NewIdent(recvName)},
				Type:  &ast.StarExpr{X: name},
			}},
		}

		funcType, ok := list.Type.(*ast.FuncType)
		if !ok {
			// TODO: When will this happen?
			continue
		}

		funcDecls = append(funcDecls, &ast.FuncDecl{
			Recv: recv,
			Name: list.Names[0],
			Type: funcType,
			Body: createBlockStmt(list.Names[0].Name, funcType),
		})
	}
	return funcDecls
}

func createBlockStmt(fname string, f *ast.FuncType) *ast.BlockStmt {
	var list []ast.Stmt

	// add called assignment statement
	calledStmt := &ast.AssignStmt{
		Lhs: []ast.Expr{
			&ast.SelectorExpr{
				X:   ast.NewIdent(recvName),
				Sel: ast.NewIdent(fname + "_Called"),
			},
		},
		Tok: token.ASSIGN,
		Rhs: []ast.Expr{ast.NewIdent("true")},
	}
	list = append(list, calledStmt)

	// add assignment for each param
	offset := 0
	for idx, field := range f.Params.List {
		// make assignments for multiple fields with same type
		if len(field.Names) > 1 {
			for _, name := range field.Names {
				assignStmt := &ast.AssignStmt{
					Lhs: []ast.Expr{
						&ast.SelectorExpr{
							X: &ast.SelectorExpr{
								X:   ast.NewIdent(recvName),
								Sel: ast.NewIdent(fname + inputSuffix),
							},
							Sel: ast.NewIdent(fmt.Sprintf("%s%d", argPrefix, idx+offset)),
						},
					},
					Tok: token.ASSIGN,
					Rhs: []ast.Expr{name},
				}
				list = append(list, assignStmt)

				offset++
			}
		} else {
			assignStmt := &ast.AssignStmt{
				Lhs: []ast.Expr{
					&ast.SelectorExpr{
						X: &ast.SelectorExpr{
							X:   ast.NewIdent(recvName),
							Sel: ast.NewIdent(fname + inputSuffix),
						},
						Sel: ast.NewIdent(fmt.Sprintf("%s%d", argPrefix, idx+offset)),
					},
				},
				Tok: token.ASSIGN,
				Rhs: []ast.Expr{field.Names[0]},
			}
			list = append(list, assignStmt)
		}
	}

	// add return statement if there are values to return
	var results []ast.Expr
	for idx := range f.Results.List {
		results = append(results, &ast.SelectorExpr{
			X: &ast.SelectorExpr{
				X:   ast.NewIdent(recvName), // for spy
				Sel: ast.NewIdent(fname + outputSuffix),
			},
			Sel: ast.NewIdent(fmt.Sprintf("%s%d", retPrefix, idx)),
		})
	}
	if len(results) > 0 {
		list = append(list, &ast.ReturnStmt{Results: results})
	}

	return &ast.BlockStmt{List: list}
}
<|endoftext|>


<commit_before>package gapi

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
)

func (c *Client) CreateUser(user User) (int64, error) {
	id := int64(0)
	data, err := json.Marshal(user)
	if err != nil {
		return id, err
	}

	req, err := c.newRequest("POST", "/api/admin/users", nil, bytes.NewBuffer(data))
	if err != nil {
		return id, err
	}
	resp, err := c.Do(req)
	if err != nil {
		return id, err
	}
	if resp.StatusCode != 200 {
		return id, errors.New(resp.Status)
	}
	data, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		return id, err
	}
	created := struct {
		Id int64 `json:"id"`
	}{}
	err = json.Unmarshal(data, &created)
	if err != nil {
		return id, err
	}
	return created.Id, err
}

func (c *Client) DeleteUser(id int64) error {
	req, err := c.newRequest("DELETE", fmt.Sprintf("/api/admin/users/%d", id), nil, nil)
	if err != nil {
		return err
	}
	resp, err := c.Do(req)
	if err != nil {
		return err
	}
	if resp.StatusCode != 200 {
		return errors.New(resp.Status)
	}
	return err
}
<commit_msg>admin methods use common request method<commit_after>package gapi

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
)

// CreateUser creates a Grafana user.
func (c *Client) CreateUser(user User) (int64, error) {
	id := int64(0)
	data, err := json.Marshal(user)
	if err != nil {
		return id, err
	}

	resp, err := c.request("POST", "/api/admin/users", nil, bytes.NewBuffer(data))
	if err != nil {
		return id, err
	}

	data, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		return id, err
	}
	created := struct {
		Id int64 `json:"id"`
	}{}
	err = json.Unmarshal(data, &created)
	if err != nil {
		return id, err
	}
	return created.Id, err
}

// DeleteUser deletes a Grafana user.
func (c *Client) DeleteUser(id int64) error {
	_, err := c.request("DELETE", fmt.Sprintf("/api/admin/users/%d", id), nil, nil)

	return err
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"math/rand"
	"net/http"
	"strconv"
	"strings"
	"time"

	openrtb "gopkg.in/bsm/openrtb.v2"
)

//OutputPerSeconds = number of seconds between stat output
var OutputPerSeconds = 10

//Wins = counter for # of wins in the last OutputPerSeconds
var Wins = 0

//Events = counter for # of events in the last OutputPerSeconds
var Events = 0

//Bids = counter for # of bids in the last OutputPerSeconds
var Bids = 0

var outputChannel = make(chan bool)

// Creative ...
type Creative struct {
	Format         string           `json:"format"`
	ID             int              `json:"id"`
	UUID           string           `json:"uuid,omitempty"`
	Name           string           `json:"name"`
	ProviderConfig *json.RawMessage `json:"providerConfig"`
}

// AgentConfig This is the agent configuration that will be sent to RTBKIT's ACS
type AgentConfig struct {
	// We use `RawMessage` for Augmentations and BidcControl, because we
	// don't need it, we just cache it.
	Account            []string         `json:"account"`
	Augmentations      *json.RawMessage `json:"augmentations"`
	BidControl         *json.RawMessage `json:"bidControl"`
	BidProbability     float64          `json:"bidProbability"`
	Creatives          []Creative       `json:"creatives"`
	ErrorFormat        string           `json:"errorFormat"`
	External           bool             `json:"external"`
	ExternalID         int              `json:"externalId"`
	LossFormat         string           `json:"lossFormat"`
	MinTimeAvailableMs float64          `json:"minTimeAvailableMs"`
	ProviderConfig     *json.RawMessage `json:"providerConfig"`
	WinFormat          string           `json:"winFormat"`
	BidderInterface    string           `json:"bidderInterface"`
}

// Agent This represents a RTBKIT Agent
type Agent struct {
	Name   string      `json:"name"`
	Config AgentConfig `json:"config"`

	// This is the price the agent will pay per impression. "Fixed price bidder".
	Price float64 `json:"price"`

	// For pacing the budgeting
	Period  int `json:"period"`
	Balance int `json:"balance"` // In microdollars

	// private state of each agent
	registered bool      // did we register the configuration in the ACS?
	pacer      chan bool // go routine updating balance in the banker
	bidID      int       // unique id for response
}

type creativesKey struct {
	// This is used to make a mapping between an impression and the
	// external-id of an agent to the creatives that can be sent to the
	// exchange for that impression.
	ImpID string
	ExtID int
}

// RegisterAgent in the ACS sending a HTTP request to the service on `acsIp`:`acsPort`
func (agent *Agent) RegisterAgent(httpClient *http.Client, acsIP string, acsPort int) {
	url := fmt.Sprintf("http://%s:%d/v1/agents/%s/config", acsIP, acsPort, agent.Name)
	body, _ := json.Marshal(agent.Config)
	reader := bytes.NewReader(body)
	req, _ := http.NewRequest("POST", url, reader)
	req.Header.Add("Accept", "application/json")
	res, err := httpClient.Do(req)
	if err != nil {
		log.Printf("ACS registration failed with %s\n", err)
		return
	}
	agent.registered = true
	res.Body.Close()
}

// UnregisterAgent Removes the agent configuration from the ACS
func (agent *Agent) UnregisterAgent(
	httpClient *http.Client, acsIP string, acsPort int) {
	url := fmt.Sprintf("http://%s:%d/v1/agents/%s/config", acsIP, acsPort, agent.Name)
	req, _ := http.NewRequest("DELETE", url, bytes.NewBufferString(""))
	res, err := httpClient.Do(req)
	if err != nil {
		log.Printf("Unregister failed with %s\n", err)
		return
	}
	agent.registered = false
	res.Body.Close()
}

func pace(httpClient *http.Client, url string, body string) {
	log.Println("Pacing...")
	req, _ := http.NewRequest("POST", url, strings.NewReader(body))
	req.Header.Add("Accept", "application/json")
	res, err := httpClient.Do(req)
	if err != nil {
		log.Printf("Balance failed with %s\n", err)
		return
	}
	res.Body.Close()
}

// StartPacer Starts a go routine which periodically updates the balance on the agents account.
func (agent *Agent) StartPacer(
	httpClient *http.Client, bankerIP string, bankerPort int) {

	accounts := agent.Config.Account

	url := fmt.Sprintf("http://%s:%d/v1/accounts/%s/balance",
		bankerIP, bankerPort, strings.Join(accounts, ":"))
	body := fmt.Sprintf("{\"USD/1M\": %d}", agent.Balance)
	ticker := time.NewTicker(time.Duration(agent.Period) * time.Millisecond)
	agent.pacer = make(chan bool)

	// Pace at the start
	go pace(httpClient, url, body)

	go func() {
		for {
			select {
			case <-ticker.C:
				// make this a new go routine?
				go pace(httpClient, url, body)
			case <-agent.pacer:
				ticker.Stop()
				return
			}
		}
	}()
}

// StartStatOutput is responsible for displaying the number of wins and events per timer tick
func StartStatOutput() {
	tickerChannel := time.NewTicker(time.Second * time.Duration(OutputPerSeconds)).C

	go func() {
		for {
			select {
			case <-tickerChannel:
				printStats()
			case <-outputChannel:
				return
			}
		}
	}()
}

// BidWin ...
func BidWin() {
	Wins++
}

// BidEvent ...
func BidEvent() {
	Events++
}

// BidIncoming ...
func BidIncoming() {
	Bids++
}

func printStats() {
	tempWins := Wins
	Wins = 0
	winsPerSecond := tempWins / OutputPerSeconds
	tempEvents := Events
	Events = 0
	eventsPerSecond := tempEvents / OutputPerSeconds
	tempBids := Bids
	Bids = 0
	bidsPerSecond := tempBids / OutputPerSeconds
	log.Println("***********************")
	log.Printf("Bids: %d (%d/second)", tempBids, bidsPerSecond)
	log.Printf("Wins: %d (%d/second)", tempWins, winsPerSecond)
	log.Printf("Events: %d (%d/second)", tempEvents, eventsPerSecond)
}

// StopPacer Stops the go routine updating the bank balance.
func (agent *Agent) StopPacer() {
	close(outputChannel)
	close(agent.pacer)
}

func round(f float64) float64 {
	return math.Floor(f + .5)
}

func random(min, max int) int {
	rand.Seed(time.Now().Unix())
	return rand.Intn(max-min) + min
}

func randomPrice(percentage float64, mainPrice float64) float64 {
	rm := (mainPrice * 100)

	diff := round(rm * percentage)

	low := int(rm - diff)
	high := int(rm + diff)

	randDec := random(low, high)
	decimal := float64(randDec) / 100.0

	return decimal
}

// DoBid Adds to the bid response the bid by the agent. The Bid is added to
// the only seat of the response. It picks a random creative from
// the list of creatives from the `Agent.Config.Creative` and places it
// in the bid.
func (agent *Agent) DoBid(
	req *openrtb.BidRequest, res *openrtb.BidResponse, ids map[creativesKey]interface{}) (*openrtb.BidResponse, bool) {

	for _, imp := range req.Imp {
		key := creativesKey{ImpID: imp.ID, ExtID: agent.Config.ExternalID}
		if ids[key] == nil {
			continue
		}

		creativeList := ids[key].([]interface{})
		// pick a random creative
		n := rand.Intn(len(creativeList))

		// Extract a usable creative ID from the JSON parse
		crid := strconv.Itoa(int(creativeList[n].(float64)))

		bidID := strconv.Itoa(agent.bidID)

		price := randomPrice(0.25, 4.00)

		log.Println("The price bid is: ", price)

		ext := map[string]interface{}{"priority": 1.0, "external-id": agent.Config.ExternalID}
		jsonExt, _ := json.Marshal(ext)
		bid := openrtb.Bid{ID: bidID, ImpID: imp.ID, CreativeID: crid, Price: price, Ext: jsonExt}
		agent.bidID++
		res.SeatBid[0].Bid = append(res.SeatBid[0].Bid, bid)
	}

	res.Currency = "USD"
	res.BidID = strconv.Itoa(agent.bidID)

	return res, len(res.SeatBid[0].Bid) > 0
}

func externalIdsFromRequest(req *openrtb.BidRequest) map[creativesKey]interface{} {
	// This function makes a mappping with a range of type (Impression Id, External Id)
	// to a slice of "creative indexes" (See the agent configuration "creative").
	// We use this auxiliary function in `DoBid` to match the `BidRequest` to the
	// creatives of the agent and create a response.
	ids := make(map[creativesKey]interface{})

	for _, imp := range req.Imp {
		log.Print("")
		var extJSON map[string]interface{}
		_ = json.Unmarshal(imp.Ext, &extJSON)

		for _, extID := range extJSON["external-ids"].([]interface{}) {
			extID = int(extID.(float64))
			key := creativesKey{ImpID: imp.ID, ExtID: extID.(int)}
			creatives := (extJSON["creative-ids"].(map[string]interface{}))[strconv.Itoa(extID.(int))]
			ids[key] = creatives.(interface{})
		}
	}
	return ids
}

func emptyResponseWithOneSeat(req *openrtb.BidRequest) *openrtb.BidResponse {
	// This function adds a Seat to the Response.
	// Seat: A buyer entity that uses a Bidder to obtain impressions on its behalf.
	seat := openrtb.SeatBid{Bid: make([]openrtb.Bid, 0)}
	seatbid := []openrtb.SeatBid{seat}
	res := &openrtb.BidResponse{ID: req.ID, SeatBid: seatbid}
	return res
}

// LoadAgent Parse a JSON file and return an Agent.
func LoadAgent(filepath string) (Agent, error) {
	var agent Agent
	data, err := ioutil.ReadFile(filepath)
	if err != nil {
		return Agent{}, err
	}
	err = json.Unmarshal(data, &agent)
	if err != nil {
		return Agent{}, err
	}
	return agent, nil
}

func loadAgents(data []byte) ([]Agent, error) {
	type Agents []Agent
	var agents Agents

	err := json.Unmarshal(data, &agents)
	if err != nil {
		return nil, err
	}
	return agents, nil
}

// LoadAgentsFromFile Parse a JSON file and return a list of Agents.
func LoadAgentsFromFile(filepath string) ([]Agent, error) {
	data, err := ioutil.ReadFile(filepath)
	if err != nil {
		return nil, err
	}
	return loadAgents(data)
}

func loadAgentsFromString(agentsString string) ([]Agent, error) {
	return loadAgents([]byte(agentsString))
}

// FindCreativeIndexFromID takes a creative ID and an AgentConfig,
// returning a creative index usable by the RTBKit router
func FindCreativeIndexFromID(crid int, agent AgentConfig) (string, error) {
	for creativeIndex, creative := range agent.Creatives {
		if creative.ID == crid {
			return strconv.Itoa(creativeIndex), nil
		}
	}
	return "", errors.New("Unable to find matching creative")
}
<commit_msg>Upped price to 10<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"math/rand"
	"net/http"
	"strconv"
	"strings"
	"time"

	openrtb "gopkg.in/bsm/openrtb.v2"
)

//OutputPerSeconds = number of seconds between stat output
var OutputPerSeconds = 10

//Wins = counter for # of wins in the last OutputPerSeconds
var Wins = 0

//Events = counter for # of events in the last OutputPerSeconds
var Events = 0

//Bids = counter for # of bids in the last OutputPerSeconds
var Bids = 0

var outputChannel = make(chan bool)

// Creative ...
type Creative struct {
	Format         string           `json:"format"`
	ID             int              `json:"id"`
	UUID           string           `json:"uuid,omitempty"`
	Name           string           `json:"name"`
	ProviderConfig *json.RawMessage `json:"providerConfig"`
}

// AgentConfig This is the agent configuration that will be sent to RTBKIT's ACS
type AgentConfig struct {
	// We use `RawMessage` for Augmentations and BidcControl, because we
	// don't need it, we just cache it.
	Account            []string         `json:"account"`
	Augmentations      *json.RawMessage `json:"augmentations"`
	BidControl         *json.RawMessage `json:"bidControl"`
	BidProbability     float64          `json:"bidProbability"`
	Creatives          []Creative       `json:"creatives"`
	ErrorFormat        string           `json:"errorFormat"`
	External           bool             `json:"external"`
	ExternalID         int              `json:"externalId"`
	LossFormat         string           `json:"lossFormat"`
	MinTimeAvailableMs float64          `json:"minTimeAvailableMs"`
	ProviderConfig     *json.RawMessage `json:"providerConfig"`
	WinFormat          string           `json:"winFormat"`
	BidderInterface    string           `json:"bidderInterface"`
}

// Agent This represents a RTBKIT Agent
type Agent struct {
	Name   string      `json:"name"`
	Config AgentConfig `json:"config"`

	// This is the price the agent will pay per impression. "Fixed price bidder".
	Price float64 `json:"price"`

	// For pacing the budgeting
	Period  int `json:"period"`
	Balance int `json:"balance"` // In microdollars

	// private state of each agent
	registered bool      // did we register the configuration in the ACS?
	pacer      chan bool // go routine updating balance in the banker
	bidID      int       // unique id for response
}

type creativesKey struct {
	// This is used to make a mapping between an impression and the
	// external-id of an agent to the creatives that can be sent to the
	// exchange for that impression.
	ImpID string
	ExtID int
}

// RegisterAgent in the ACS sending a HTTP request to the service on `acsIp`:`acsPort`
func (agent *Agent) RegisterAgent(httpClient *http.Client, acsIP string, acsPort int) {
	url := fmt.Sprintf("http://%s:%d/v1/agents/%s/config", acsIP, acsPort, agent.Name)
	body, _ := json.Marshal(agent.Config)
	reader := bytes.NewReader(body)
	req, _ := http.NewRequest("POST", url, reader)
	req.Header.Add("Accept", "application/json")
	res, err := httpClient.Do(req)
	if err != nil {
		log.Printf("ACS registration failed with %s\n", err)
		return
	}
	agent.registered = true
	res.Body.Close()
}

// UnregisterAgent Removes the agent configuration from the ACS
func (agent *Agent) UnregisterAgent(
	httpClient *http.Client, acsIP string, acsPort int) {
	url := fmt.Sprintf("http://%s:%d/v1/agents/%s/config", acsIP, acsPort, agent.Name)
	req, _ := http.NewRequest("DELETE", url, bytes.NewBufferString(""))
	res, err := httpClient.Do(req)
	if err != nil {
		log.Printf("Unregister failed with %s\n", err)
		return
	}
	agent.registered = false
	res.Body.Close()
}

func pace(httpClient *http.Client, url string, body string) {
	log.Println("Pacing...")
	req, _ := http.NewRequest("POST", url, strings.NewReader(body))
	req.Header.Add("Accept", "application/json")
	res, err := httpClient.Do(req)
	if err != nil {
		log.Printf("Balance failed with %s\n", err)
		return
	}
	res.Body.Close()
}

// StartPacer Starts a go routine which periodically updates the balance on the agents account.
func (agent *Agent) StartPacer(
	httpClient *http.Client, bankerIP string, bankerPort int) {

	accounts := agent.Config.Account

	url := fmt.Sprintf("http://%s:%d/v1/accounts/%s/balance",
		bankerIP, bankerPort, strings.Join(accounts, ":"))
	body := fmt.Sprintf("{\"USD/1M\": %d}", agent.Balance)
	ticker := time.NewTicker(time.Duration(agent.Period) * time.Millisecond)
	agent.pacer = make(chan bool)

	// Pace at the start
	go pace(httpClient, url, body)

	go func() {
		for {
			select {
			case <-ticker.C:
				// make this a new go routine?
				go pace(httpClient, url, body)
			case <-agent.pacer:
				ticker.Stop()
				return
			}
		}
	}()
}

// StartStatOutput is responsible for displaying the number of wins and events per timer tick
func StartStatOutput() {
	tickerChannel := time.NewTicker(time.Second * time.Duration(OutputPerSeconds)).C

	go func() {
		for {
			select {
			case <-tickerChannel:
				printStats()
			case <-outputChannel:
				return
			}
		}
	}()
}

// BidWin ...
func BidWin() {
	Wins++
}

// BidEvent ...
func BidEvent() {
	Events++
}

// BidIncoming ...
func BidIncoming() {
	Bids++
}

func printStats() {
	tempWins := Wins
	Wins = 0
	winsPerSecond := tempWins / OutputPerSeconds
	tempEvents := Events
	Events = 0
	eventsPerSecond := tempEvents / OutputPerSeconds
	tempBids := Bids
	Bids = 0
	bidsPerSecond := tempBids / OutputPerSeconds
	log.Println("***********************")
	log.Printf("Bids: %d (%d/second)", tempBids, bidsPerSecond)
	log.Printf("Wins: %d (%d/second)", tempWins, winsPerSecond)
	log.Printf("Events: %d (%d/second)", tempEvents, eventsPerSecond)
}

// StopPacer Stops the go routine updating the bank balance.
func (agent *Agent) StopPacer() {
	close(outputChannel)
	close(agent.pacer)
}

func round(f float64) float64 {
	return math.Floor(f + .5)
}

func random(min, max int) int {
	rand.Seed(time.Now().Unix())
	return rand.Intn(max-min) + min
}

func randomPrice(percentage float64, mainPrice float64) float64 {
	rm := (mainPrice * 100)

	diff := round(rm * percentage)

	low := int(rm - diff)
	high := int(rm + diff)

	randDec := random(low, high)
	decimal := float64(randDec) / 100.0

	return decimal
}

// DoBid Adds to the bid response the bid by the agent. The Bid is added to
// the only seat of the response. It picks a random creative from
// the list of creatives from the `Agent.Config.Creative` and places it
// in the bid.
func (agent *Agent) DoBid(
	req *openrtb.BidRequest, res *openrtb.BidResponse, ids map[creativesKey]interface{}) (*openrtb.BidResponse, bool) {

	for _, imp := range req.Imp {
		key := creativesKey{ImpID: imp.ID, ExtID: agent.Config.ExternalID}
		if ids[key] == nil {
			continue
		}

		creativeList := ids[key].([]interface{})
		// pick a random creative
		n := rand.Intn(len(creativeList))

		// Extract a usable creative ID from the JSON parse
		crid := strconv.Itoa(int(creativeList[n].(float64)))

		bidID := strconv.Itoa(agent.bidID)

		price := randomPrice(0.25, 10.00)

		log.Println("The price bid is: ", price)

		ext := map[string]interface{}{"priority": 1.0, "external-id": agent.Config.ExternalID}
		jsonExt, _ := json.Marshal(ext)
		bid := openrtb.Bid{ID: bidID, ImpID: imp.ID, CreativeID: crid, Price: price, Ext: jsonExt}
		agent.bidID++
		res.SeatBid[0].Bid = append(res.SeatBid[0].Bid, bid)
	}

	res.Currency = "USD"
	res.BidID = strconv.Itoa(agent.bidID)

	return res, len(res.SeatBid[0].Bid) > 0
}

func externalIdsFromRequest(req *openrtb.BidRequest) map[creativesKey]interface{} {
	// This function makes a mappping with a range of type (Impression Id, External Id)
	// to a slice of "creative indexes" (See the agent configuration "creative").
	// We use this auxiliary function in `DoBid` to match the `BidRequest` to the
	// creatives of the agent and create a response.
	ids := make(map[creativesKey]interface{})

	for _, imp := range req.Imp {
		log.Print("")
		var extJSON map[string]interface{}
		_ = json.Unmarshal(imp.Ext, &extJSON)

		for _, extID := range extJSON["external-ids"].([]interface{}) {
			extID = int(extID.(float64))
			key := creativesKey{ImpID: imp.ID, ExtID: extID.(int)}
			creatives := (extJSON["creative-ids"].(map[string]interface{}))[strconv.Itoa(extID.(int))]
			ids[key] = creatives.(interface{})
		}
	}
	return ids
}

func emptyResponseWithOneSeat(req *openrtb.BidRequest) *openrtb.BidResponse {
	// This function adds a Seat to the Response.
	// Seat: A buyer entity that uses a Bidder to obtain impressions on its behalf.
	seat := openrtb.SeatBid{Bid: make([]openrtb.Bid, 0)}
	seatbid := []openrtb.SeatBid{seat}
	res := &openrtb.BidResponse{ID: req.ID, SeatBid: seatbid}
	return res
}

// LoadAgent Parse a JSON file and return an Agent.
func LoadAgent(filepath string) (Agent, error) {
	var agent Agent
	data, err := ioutil.ReadFile(filepath)
	if err != nil {
		return Agent{}, err
	}
	err = json.Unmarshal(data, &agent)
	if err != nil {
		return Agent{}, err
	}
	return agent, nil
}

func loadAgents(data []byte) ([]Agent, error) {
	type Agents []Agent
	var agents Agents

	err := json.Unmarshal(data, &agents)
	if err != nil {
		return nil, err
	}
	return agents, nil
}

// LoadAgentsFromFile Parse a JSON file and return a list of Agents.
func LoadAgentsFromFile(filepath string) ([]Agent, error) {
	data, err := ioutil.ReadFile(filepath)
	if err != nil {
		return nil, err
	}
	return loadAgents(data)
}

func loadAgentsFromString(agentsString string) ([]Agent, error) {
	return loadAgents([]byte(agentsString))
}

// FindCreativeIndexFromID takes a creative ID and an AgentConfig,
// returning a creative index usable by the RTBKit router
func FindCreativeIndexFromID(crid int, agent AgentConfig) (string, error) {
	for creativeIndex, creative := range agent.Creatives {
		if creative.ID == crid {
			return strconv.Itoa(creativeIndex), nil
		}
	}
	return "", errors.New("Unable to find matching creative")
}
<|endoftext|>


<commit_before>package skunk

import (
	"bytes"
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"sort"
	"text/tabwriter"
	"time"
)

// MinuteCycle, QuarterHourCycle, HalfHourCycle, and HourCycle all represent useful reporting cycles for an agent. Other
// cycles are permitted, provide they would not violate the two-POSTs-per-minute limit on NewRelic APIs.
const (
	MinuteCycle   = time.Minute
	HalfHourCycle = time.Minute * 30
	HourCycle     = time.Hour
)

// NewRelicAPI is the URL to POST NewRelic metrics data to. This may be altered to change the
// default endpoint of new agents. Already-initialized agents do not use this.
var NewRelicAPI = `https://platform-api.newrelic.com/platform/v1/metrics`

type opFunc func(*Agent) error

type Agent struct {
	// Initialization fields -- these may not change after Start is called. Prior to calling Start, you may tweak
	// them to your heart's content.
	Cycle      time.Duration
	Client     *http.Client
	Log        io.Writer
	LogMetrics bool

	apiURL string
	apiKey string

	// Access to the following fields controlled by runloop after init
	body     *Body
	err      error
	lastPoll time.Time
	ticker   *time.Ticker
	ops      chan<- opFunc
}

func New(version, apiKey string) (*Agent, error) {
	host, err := os.Hostname()
	if err != nil {
		// Hostnames are required.
		return nil, err
	}

	rep := AgentRep{
		Host:    host,
		PID:     os.Getpid(),
		Version: version,
	}

	return NewWithRep(apiKey, rep)
}

func NewWithRep(apiKey string, rep AgentRep) (agent *Agent, err error) {
	switch {
	case len(apiKey) == 0:
		return nil, mkerr(ErrNoAPIKey, nil)
	case len(rep.Version) == 0:
		return nil, mkerr(ErrNoVersion, nil)
	case len(rep.Host) == 0:
		return nil, mkerr(ErrNoHost, nil)
	}

	return &Agent{
		Client:     http.DefaultClient,
		Cycle:      MinuteCycle,
		Log:        ioutil.Discard,
		LogMetrics: false,

		apiURL: NewRelicAPI,
		apiKey: apiKey,

		body: &Body{
			Agent:      rep,
			Components: nil,
		},

		lastPoll: time.Now(),
		ticker:   nil,
	}, nil
}

func (a *Agent) Start() {
	if a.ops != nil {
		return
	}

	if a.Log == nil {
		a.Log = ioutil.Discard
	}

	ops := make(chan opFunc)
	a.ops = ops

	go a.run(ops)
}

// Close kills the agent's runloop and makes it completely inert. Using the agent afterward will result in a panic. Any
// error held by the agent prior to shutdown is returned.
//
// When calling Close, you must ensure that the agent is no longer in use and will not be used by any goroutine after
// Close is called.
func (a *Agent) Close() error {
	err := a.Err()
	a.ops <- shutdown
	return err
}

// shutdown is an opFunc that closes an agent's ops channel.
func shutdown(a *Agent) error {
	if err := a.sendRequest(time.Now()); iserr(err, errMustRetry) {
		fmt.Fprintln(a.Log, "skunk: received 50x error from NewRelic on shutdown flush - dropping payload on the floor")
	} else if err != nil {
		a.err = err
	}
	close(a.ops)
	return mkerr(errShuttingDown, nil)
}

type opGetErr chan<- error

func (c opGetErr) Exec(a *Agent) error {
	c <- a.err
	return nil
}

func (a *Agent) Err() (err error) {
	out := make(chan error)
	a.ops <- opGetErr(out).Exec
	return <-out
}

func (a *Agent) run(ops <-chan opFunc) {
	var timer *time.Timer
	var retry <-chan time.Time
	retryNeeded := false

	a.ticker = time.NewTicker(a.Cycle)
	defer a.ticker.Stop()

	trySend := func(from time.Time) {
		if err := a.sendRequest(from); err == nil {
			retryNeeded = false
			a.lastPoll = from
			a.clear()
		} else if iserr(err, errMustRetry) {
			if timer == nil {
				timer = time.NewTimer(time.Minute)
				retry = timer.C
			} else {
				timer.Reset(time.Minute)
			}
			retryNeeded = true
		} else {
			a.err = err
		}
	}

	for {
		select {
		case from := <-retry:
			trySend(from)
		case from := <-a.ticker.C:
			if a.LogMetrics {
				a.logMetrics()
			}

			if !retryNeeded {
				// Let the retry loop take over until things are back to normal.
				trySend(from)
			}
		case op, ok := <-ops:
			if !ok {
				return
			} else if op == nil {
				// This should be impossible. If it happens, log it and skip the op.
				fmt.Fprintln(a.Log, ErrNilOpReceived)
				continue
			}

			if err := op(a); iserr(err, errShuttingDown) {
				return
			} else if err != nil {
				a.err = err
			}
		}
	}
}

func (a *Agent) logMetrics() {
	if a.Log == ioutil.Discard {
		return
	}

	components := make([]Component, len(a.body.Components))
	i := 0
	for _, c := range a.body.Components {
		if len(c.Metrics) > 0 {
			components[i] = *c
			i++
		}
	}

	if i == 0 {
		return
	}

	go logComponentMetrics(a.Log, components[:i])
}

func logComponentMetrics(w io.Writer, components []Component) {
	if len(components) == 0 {
		return
	}

	var buf bytes.Buffer
	tw := tabwriter.NewWriter(&buf, 1, 8, 1, ' ', tabwriter.TabIndent)

	var keys []string
	for _, com := range components {
		if len(com.Metrics) == 0 {
			continue
		}

		// Sort keys so the output is easier to sift through.
		for key := range com.Metrics {
			keys = append(keys, key)
		}
		sort.Strings(keys)

		fmt.Fprintf(tw, "%s (%s) %v\n\tName\tCount\tTotal\tAverage\tMin\tMax\tSS\n", com.Name, com.GUID, com.Duration.Duration)
		for _, key := range keys {
			m := com.Metrics[key]
			switch m := m.(type) {
			case RangeMetric:
				fmt.Fprintf(tw, "\t%s\t%v\t%v\t%v\t%v\t%v\t%v\n",
					key, m.Count, m.Total, m.Total/float64(m.Count), m.Min, m.Max, m.Square-((m.Total*m.Total)/float64(m.Count)))
			case ScalarMetric:
				f := float64(m)
				fmt.Fprintf(tw, "\t%s\t1\t%v\t%v\t%v\t%v\t0\n",
					key, f, f, f, f)
			}
		}
		tw.Flush()

		keys = keys[0:0]
	}

	if _, err := buf.WriteTo(w); err != nil {
		fmt.Fprintf(w, "skunk: error writing metrics log entries: %v", err)
	}
}

func (a *Agent) sendRequest(from time.Time) (err error) {
	var buf bytes.Buffer
	compressed := true
tryGetPayload:
	err = a.getPayload(&buf, from, compressed)
	switch {
	case err == nil:
	case iserr(err, errNoMetrics):
		return nil // Nothing to do.
	default:
		if _, ok := err.(*json.MarshalerError); ok {
			// Can't do anything about this. This error might be worth panicking over.
			return mkerr(ErrEncodingJSON, err)
		}

		if compressed {
			// Try without compression in case it's some anomalous unknown compression error that's eluded
			// everyone but me (i.e., should be almost impossible).
			compressed = false
			buf.Reset()
			goto tryGetPayload
		}
		return err
	}

	req, err := http.NewRequest("POST", a.apiURL, &buf)
	if err != nil {
		// No idea what happened here, assume the worst.
		return err
	}

	// Set headers
	req.Header.Set("X-License-Key", a.apiKey)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")
	if compressed {
		req.Header.Set("Content-Encoding", "gzip")
	}

	resp, err := a.Client.Do(req)
	if resp != nil {
		defer func() {
			closeErr := resp.Body.Close()
			if closeErr != nil {
				fmt.Fprintf(a.Log, "skunk: error closing response body: %v", closeErr)
			}
		}()
	}

	if err != nil {
		return err
	}

	if resp.StatusCode == 200 {
		return nil
	}

	var nrErr struct {
		Error string `json:"error"`
	}
	decoder := json.NewDecoder(resp.Body)
	if err = decoder.Decode(nrErr); err == nil && len(nrErr.Error) > 0 {
		fmt.Fprintf(a.Log, "skunk: received NewRelic error: %s", nrErr.Error)
	}

	return statusError(resp)
}

func statusError(resp *http.Response) error {
	code := resp.StatusCode
	switch {
	case code >= 200 && code < 300:
		return nil
	case code == 400:
		return mkerr(ErrBadPayload, nil)
	case code == 403:
		return mkerr(ErrForbidden, nil)
	case code == 404:
		return mkerr(ErrBadRequest, nil)
	case code == 405:
		return mkerr(ErrBadRequest, nil)
	case code == 413:
		return mkerr(ErrBodyTooLarge, nil)
	case code >= 500 && code < 600:
		return mkerr(errMustRetry, nil)
	default:
		return fmt.Errorf("skunk: got unexpected status code %d %s from NewRelic.", code, resp.Status)
	}
}

// Component gets a component with the given name and GUID from the Agent. If no such component exists, then a new one
// is allocated and it is returned. No tests are done to ensure that components with the same name but a different GUID
// or vice versa are allocated, so it is possible to end up with potentially inconsistent data.
func (a *Agent) Component(name, guid string) (*Component, error) {
	if len(name) > 32 {
		// Per NewRelic, names must be <= 32 characters in length.
		return nil, mkerr(ErrNameTooLong, nil)
	}

	out := make(chan *Component)
	a.ops <- (addComponent{name, guid, out}).Exec
	return <-out, nil
}

// addComponent is a small wrapper around a parameter bundle to create a component and add it to an agent, or raturn an
// agent with the same name and guid.
type addComponent struct {
	name, guid string
	out        chan<- *Component
}

// Exec searches for a component of the name and guid described by ac. If it finds a component in the agent, a, it is
// sent on ac's out channel. Otherwise, a new component is created, added to the agent, and sent on the out channel.
func (ac addComponent) Exec(a *Agent) error {
	body := a.body
	for _, c := range body.Components {
		if c.Name == ac.name && c.GUID == ac.guid {
			ac.out <- c
			return nil
		}
	}

	c := &Component{
		Name:    ac.name,
		GUID:    ac.guid,
		Metrics: make(map[string]Metric),
		agent:   a,
	}
	body.Components = append(body.Components, c)
	ac.out <- c

	return nil
}

// clear empties out all metrics held by the agent. This should be called after a payload has been successfully sent to
// reset all components to a pristine state.
func (a *Agent) clear() {
	for _, c := range a.body.Components {
		// Allocate a new Metrics map and reuse the old map's length as its capacity.
		c.Metrics = make(map[string]Metric, len(c.Metrics))
		c.Duration.Duration = 0 // Should be zero, but clear() says it'll make it pristine, so zero it anyway.
		c.start = time.Time{}
	}
}

// getPayload returns a JSON payload as a byte slice to send to NewRelic as its POSTed body. The resulting JSON does not
// include components without metrics.
func (a *Agent) getPayload(w io.Writer, from time.Time, compressed bool) (err error) {
	// Make a copy of the body and exclude components without metrics.
	body := *a.body
	body.Components = make([]*Component, 0, len(body.Components))
	for _, com := range a.body.Components {
		if len(com.Metrics) == 0 || com.start.IsZero() {
			continue
		}

		dupe := *com
		dupe.Duration.Duration = from.Sub(com.start)
		if dupe.Duration.Duration < 0 {
			// Metrics from the future aren't allowed.
			dupe.Duration.Duration = 0
		}

		body.Components = append(body.Components, &dupe)
	}

	if len(body.Components) == 0 {
		return mkerr(errNoMetrics, nil)
	}

	if compressed {
		zipWriter := gzip.NewWriter(w)
		defer func() {
			if err == nil {
				err = zipWriter.Close()
			}
		}()
		w = zipWriter
	}
	encoder := json.NewEncoder(w)
	return encoder.Encode(body)
}
<commit_msg>Add newlines to log messages<commit_after>package skunk

import (
	"bytes"
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"sort"
	"text/tabwriter"
	"time"
)

// MinuteCycle, QuarterHourCycle, HalfHourCycle, and HourCycle all represent useful reporting cycles for an agent. Other
// cycles are permitted, provide they would not violate the two-POSTs-per-minute limit on NewRelic APIs.
const (
	MinuteCycle   = time.Minute
	HalfHourCycle = time.Minute * 30
	HourCycle     = time.Hour
)

// NewRelicAPI is the URL to POST NewRelic metrics data to. This may be altered to change the
// default endpoint of new agents. Already-initialized agents do not use this.
var NewRelicAPI = `https://platform-api.newrelic.com/platform/v1/metrics`

type opFunc func(*Agent) error

type Agent struct {
	// Initialization fields -- these may not change after Start is called. Prior to calling Start, you may tweak
	// them to your heart's content.
	Cycle      time.Duration
	Client     *http.Client
	Log        io.Writer
	LogMetrics bool

	apiURL string
	apiKey string

	// Access to the following fields controlled by runloop after init
	body     *Body
	err      error
	lastPoll time.Time
	ticker   *time.Ticker
	ops      chan<- opFunc
}

func New(version, apiKey string) (*Agent, error) {
	host, err := os.Hostname()
	if err != nil {
		// Hostnames are required.
		return nil, err
	}

	rep := AgentRep{
		Host:    host,
		PID:     os.Getpid(),
		Version: version,
	}

	return NewWithRep(apiKey, rep)
}

func NewWithRep(apiKey string, rep AgentRep) (agent *Agent, err error) {
	switch {
	case len(apiKey) == 0:
		return nil, mkerr(ErrNoAPIKey, nil)
	case len(rep.Version) == 0:
		return nil, mkerr(ErrNoVersion, nil)
	case len(rep.Host) == 0:
		return nil, mkerr(ErrNoHost, nil)
	}

	return &Agent{
		Client:     http.DefaultClient,
		Cycle:      MinuteCycle,
		Log:        ioutil.Discard,
		LogMetrics: false,

		apiURL: NewRelicAPI,
		apiKey: apiKey,

		body: &Body{
			Agent:      rep,
			Components: nil,
		},

		lastPoll: time.Now(),
		ticker:   nil,
	}, nil
}

func (a *Agent) Start() {
	if a.ops != nil {
		return
	}

	if a.Log == nil {
		a.Log = ioutil.Discard
	}

	ops := make(chan opFunc)
	a.ops = ops

	go a.run(ops)
}

// Close kills the agent's runloop and makes it completely inert. Using the agent afterward will result in a panic. Any
// error held by the agent prior to shutdown is returned.
//
// When calling Close, you must ensure that the agent is no longer in use and will not be used by any goroutine after
// Close is called.
func (a *Agent) Close() error {
	err := a.Err()
	a.ops <- shutdown
	return err
}

// shutdown is an opFunc that closes an agent's ops channel.
func shutdown(a *Agent) error {
	if err := a.sendRequest(time.Now()); iserr(err, errMustRetry) {
		fmt.Fprintln(a.Log, "skunk: received 50x error from NewRelic on shutdown flush - dropping payload on the floor")
	} else if err != nil {
		a.err = err
	}
	close(a.ops)
	return mkerr(errShuttingDown, nil)
}

type opGetErr chan<- error

func (c opGetErr) Exec(a *Agent) error {
	c <- a.err
	return nil
}

func (a *Agent) Err() (err error) {
	out := make(chan error)
	a.ops <- opGetErr(out).Exec
	return <-out
}

func (a *Agent) run(ops <-chan opFunc) {
	var timer *time.Timer
	var retry <-chan time.Time
	retryNeeded := false

	a.ticker = time.NewTicker(a.Cycle)
	defer a.ticker.Stop()

	trySend := func(from time.Time) {
		if err := a.sendRequest(from); err == nil {
			retryNeeded = false
			a.lastPoll = from
			a.clear()
		} else if iserr(err, errMustRetry) {
			if timer == nil {
				timer = time.NewTimer(time.Minute)
				retry = timer.C
			} else {
				timer.Reset(time.Minute)
			}
			retryNeeded = true
		} else {
			a.err = err
		}
	}

	for {
		select {
		case from := <-retry:
			trySend(from)
		case from := <-a.ticker.C:
			if a.LogMetrics {
				a.logMetrics()
			}

			if !retryNeeded {
				// Let the retry loop take over until things are back to normal.
				trySend(from)
			}
		case op, ok := <-ops:
			if !ok {
				return
			} else if op == nil {
				// This should be impossible. If it happens, log it and skip the op.
				fmt.Fprintln(a.Log, ErrNilOpReceived)
				continue
			}

			if err := op(a); iserr(err, errShuttingDown) {
				return
			} else if err != nil {
				a.err = err
			}
		}
	}
}

func (a *Agent) logMetrics() {
	if a.Log == ioutil.Discard {
		return
	}

	components := make([]Component, len(a.body.Components))
	i := 0
	for _, c := range a.body.Components {
		if len(c.Metrics) > 0 {
			components[i] = *c
			i++
		}
	}

	if i == 0 {
		return
	}

	go logComponentMetrics(a.Log, components[:i])
}

func logComponentMetrics(w io.Writer, components []Component) {
	if len(components) == 0 {
		return
	}

	var buf bytes.Buffer
	tw := tabwriter.NewWriter(&buf, 1, 8, 1, ' ', tabwriter.TabIndent)

	var keys []string
	for _, com := range components {
		if len(com.Metrics) == 0 {
			continue
		}

		// Sort keys so the output is easier to sift through.
		for key := range com.Metrics {
			keys = append(keys, key)
		}
		sort.Strings(keys)

		fmt.Fprintf(tw, "%s (%s) %v\n\tName\tCount\tTotal\tAverage\tMin\tMax\tSS\n", com.Name, com.GUID, com.Duration.Duration)
		for _, key := range keys {
			m := com.Metrics[key]
			switch m := m.(type) {
			case RangeMetric:
				fmt.Fprintf(tw, "\t%s\t%v\t%v\t%v\t%v\t%v\t%v\n",
					key, m.Count, m.Total, m.Total/float64(m.Count), m.Min, m.Max, m.Square-((m.Total*m.Total)/float64(m.Count)))
			case ScalarMetric:
				f := float64(m)
				fmt.Fprintf(tw, "\t%s\t1\t%v\t%v\t%v\t%v\t0\n",
					key, f, f, f, f)
			}
		}
		tw.Flush()

		keys = keys[0:0]
	}

	if _, err := buf.WriteTo(w); err != nil {
		fmt.Fprintf(w, "skunk: error writing metrics log entries: %v\n", err)
	}
}

func (a *Agent) sendRequest(from time.Time) (err error) {
	var buf bytes.Buffer
	compressed := true
tryGetPayload:
	err = a.getPayload(&buf, from, compressed)
	switch {
	case err == nil:
	case iserr(err, errNoMetrics):
		return nil // Nothing to do.
	default:
		if _, ok := err.(*json.MarshalerError); ok {
			// Can't do anything about this. This error might be worth panicking over.
			return mkerr(ErrEncodingJSON, err)
		}

		if compressed {
			// Try without compression in case it's some anomalous unknown compression error that's eluded
			// everyone but me (i.e., should be almost impossible).
			compressed = false
			buf.Reset()
			goto tryGetPayload
		}
		return err
	}

	req, err := http.NewRequest("POST", a.apiURL, &buf)
	if err != nil {
		// No idea what happened here, assume the worst.
		return err
	}

	// Set headers
	req.Header.Set("X-License-Key", a.apiKey)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")
	if compressed {
		req.Header.Set("Content-Encoding", "gzip")
	}

	resp, err := a.Client.Do(req)
	if resp != nil {
		defer func() {
			closeErr := resp.Body.Close()
			if closeErr != nil {
				fmt.Fprintf(a.Log, "skunk: error closing response body: %v\n", closeErr)
			}
		}()
	}

	if err != nil {
		return err
	}

	if resp.StatusCode == 200 {
		return nil
	}

	var nrErr struct {
		Error string `json:"error"`
	}
	decoder := json.NewDecoder(resp.Body)
	if err = decoder.Decode(nrErr); err == nil && len(nrErr.Error) > 0 {
		fmt.Fprintf(a.Log, "skunk: received NewRelic error: %s\n", nrErr.Error)
	}
	if _, err := io.Copy(ioutil.Discard, resp.Body); err != nil {
		fmt.Fprintf(a.Log, "skunk: error discarding body remainder: %v\n", err)
	}

	return statusError(resp)
}

func statusError(resp *http.Response) error {
	code := resp.StatusCode
	switch {
	case code >= 200 && code < 300:
		return nil
	case code == 400:
		return mkerr(ErrBadPayload, nil)
	case code == 403:
		return mkerr(ErrForbidden, nil)
	case code == 404:
		return mkerr(ErrBadRequest, nil)
	case code == 405:
		return mkerr(ErrBadRequest, nil)
	case code == 413:
		return mkerr(ErrBodyTooLarge, nil)
	case code >= 500 && code < 600:
		return mkerr(errMustRetry, nil)
	default:
		return fmt.Errorf("skunk: got unexpected status code %d %s from NewRelic.", code, resp.Status)
	}
}

// Component gets a component with the given name and GUID from the Agent. If no such component exists, then a new one
// is allocated and it is returned. No tests are done to ensure that components with the same name but a different GUID
// or vice versa are allocated, so it is possible to end up with potentially inconsistent data.
func (a *Agent) Component(name, guid string) (*Component, error) {
	if len(name) > 32 {
		// Per NewRelic, names must be <= 32 characters in length.
		return nil, mkerr(ErrNameTooLong, nil)
	}

	out := make(chan *Component)
	a.ops <- (addComponent{name, guid, out}).Exec
	return <-out, nil
}

// addComponent is a small wrapper around a parameter bundle to create a component and add it to an agent, or raturn an
// agent with the same name and guid.
type addComponent struct {
	name, guid string
	out        chan<- *Component
}

// Exec searches for a component of the name and guid described by ac. If it finds a component in the agent, a, it is
// sent on ac's out channel. Otherwise, a new component is created, added to the agent, and sent on the out channel.
func (ac addComponent) Exec(a *Agent) error {
	body := a.body
	for _, c := range body.Components {
		if c.Name == ac.name && c.GUID == ac.guid {
			ac.out <- c
			return nil
		}
	}

	c := &Component{
		Name:    ac.name,
		GUID:    ac.guid,
		Metrics: make(map[string]Metric),
		agent:   a,
	}
	body.Components = append(body.Components, c)
	ac.out <- c

	return nil
}

// clear empties out all metrics held by the agent. This should be called after a payload has been successfully sent to
// reset all components to a pristine state.
func (a *Agent) clear() {
	for _, c := range a.body.Components {
		// Allocate a new Metrics map and reuse the old map's length as its capacity.
		c.Metrics = make(map[string]Metric, len(c.Metrics))
		c.Duration.Duration = 0 // Should be zero, but clear() says it'll make it pristine, so zero it anyway.
		c.start = time.Time{}
	}
}

// getPayload returns a JSON payload as a byte slice to send to NewRelic as its POSTed body. The resulting JSON does not
// include components without metrics.
func (a *Agent) getPayload(w io.Writer, from time.Time, compressed bool) (err error) {
	// Make a copy of the body and exclude components without metrics.
	body := *a.body
	body.Components = make([]*Component, 0, len(body.Components))
	for _, com := range a.body.Components {
		if len(com.Metrics) == 0 || com.start.IsZero() {
			continue
		}

		dupe := *com
		dupe.Duration.Duration = from.Sub(com.start)
		if dupe.Duration.Duration < 0 {
			// Metrics from the future aren't allowed.
			dupe.Duration.Duration = 0
		}

		body.Components = append(body.Components, &dupe)
	}

	if len(body.Components) == 0 {
		return mkerr(errNoMetrics, nil)
	}

	if compressed {
		zipWriter := gzip.NewWriter(w)
		defer func() {
			if err == nil {
				err = zipWriter.Close()
			}
		}()
		w = zipWriter
	}
	encoder := json.NewEncoder(w)
	return encoder.Encode(body)
}
<|endoftext|>