<commit_before>// +build e2e

/*
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 e2e

import (
	"fmt"
	"strconv"
	"testing"
	"time"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"knative.dev/serving/pkg/apis/autoscaling"
	v1 "knative.dev/serving/pkg/apis/serving/v1"
	"knative.dev/serving/pkg/resources"
	rtesting "knative.dev/serving/pkg/testing/v1"
	"knative.dev/serving/test"
	v1test "knative.dev/serving/test/v1"
)

func TestMinScale(t *testing.T) {
	t.Parallel()

	const minScale = 4

	clients := Setup(t)

	names := test.ResourceNames{
		// Config and Route have different names to avoid false positives
		Config: test.ObjectNameForTest(t),
		Route:  test.ObjectNameForTest(t),
		Image:  "helloworld",
	}

	test.EnsureTearDown(t, clients, &names)

	t.Log("Creating route")
	if _, err := v1test.CreateRoute(t, clients, names); err != nil {
		t.Fatal("Failed to create Route:", err)
	}

	t.Log("Creating configuration")
	cfg, err := v1test.CreateConfiguration(t, clients, names, withMinScale(minScale))
	if err != nil {
		t.Fatal("Failed to create Configuration:", err)
	}

	revName := latestRevisionName(t, clients, names.Config, "")
	serviceName := privateServiceName(t, clients, revName)

	t.Log("Waiting for revision to scale to minScale before becoming ready")
	if lr, err := waitForDesiredScale(clients, serviceName, gte(minScale)); err != nil {
		t.Fatalf("The revision %q scaled to %d < %d before becoming ready: %v", revName, lr, minScale, err)
	}

	t.Log("Waiting for revision to become ready")
	if err := v1test.WaitForRevisionState(
		clients.ServingClient, revName, v1test.IsRevisionReady, "RevisionIsReady",
	); err != nil {
		t.Fatalf("The Revision %q did not become ready: %v", revName, err)
	}

	t.Log("Holding revision at minScale after becoming ready")
	if lr, ok := ensureDesiredScale(clients, t, serviceName, gte(minScale)); !ok {
		t.Fatalf("The revision %q observed scale %d < %d after becoming ready", revName, lr, minScale)
	}

	t.Log("Updating configuration")
	if _, err := v1test.PatchConfig(t, clients, cfg, rtesting.WithConfigEnv(corev1.EnvVar{Name: "FOO", Value: "BAR"})); err != nil {
		t.Fatal("Failed to update Configuration:", err)
	}

	newRevName := latestRevisionName(t, clients, names.Config, revName)
	newServiceName := privateServiceName(t, clients, newRevName)

	t.Log("Waiting for new revision to scale to minScale after update")
	if lr, err := waitForDesiredScale(clients, newServiceName, gte(minScale)); err != nil {
		t.Fatalf("The revision %q scaled to %d < %d after creating route: %v", newRevName, lr, minScale, err)
	}

	t.Log("Waiting for new revision to become ready")
	if err := v1test.WaitForRevisionState(
		clients.ServingClient, newRevName, v1test.IsRevisionReady, "RevisionIsReady",
	); err != nil {
		t.Fatalf("The Revision %q did not become ready: %v", newRevName, err)
	}

	t.Log("Holding new revision at minScale after becoming ready")
	if lr, ok := ensureDesiredScale(clients, t, newServiceName, gte(minScale)); !ok {
		t.Fatalf("The revision %q observed scale %d < %d after becoming ready", newRevName, lr, minScale)
	}

	t.Log("Waiting for old revision to scale below minScale after being replaced")
	if lr, err := waitForDesiredScale(clients, serviceName, lt(minScale)); err != nil {
		t.Fatalf("The revision %q scaled to %d > %d after not being routable anymore: %v", revName, lr, minScale, err)
	}

	t.Log("Deleting route", names.Route)
	if err := clients.ServingClient.Routes.Delete(names.Route, &metav1.DeleteOptions{}); err != nil {
		t.Fatalf("Failed to delete route %q: %v", names.Route, err)
	}

	t.Log("Waiting for new revision to scale below minScale when there is no route")
	if lr, err := waitForDesiredScale(clients, newServiceName, lt(minScale)); err != nil {
		t.Fatalf("The revision %q scaled to %d > %d after not being routable anymore: %v", newRevName, lr, minScale, err)
	}
}

func gte(m int) func(int) bool {
	return func(n int) bool {
		return n >= m
	}
}

func lt(m int) func(int) bool {
	return func(n int) bool {
		return n < m
	}
}

func withMinScale(minScale int) func(cfg *v1.Configuration) {
	return func(cfg *v1.Configuration) {
		if cfg.Spec.Template.Annotations == nil {
			cfg.Spec.Template.Annotations = make(map[string]string, 1)
		}
		cfg.Spec.Template.Annotations[autoscaling.MinScaleAnnotationKey] = strconv.Itoa(minScale)
	}
}

func latestRevisionName(t *testing.T, clients *test.Clients, configName, oldRevName string) string {
	// Wait for the Config have a LatestCreatedRevisionName
	if err := v1test.WaitForConfigurationState(
		clients.ServingClient, configName,
		func(c *v1.Configuration) (bool, error) {
			return c.Status.LatestCreatedRevisionName != oldRevName, nil
		}, "ConfigurationHasUpdatedCreatedRevision",
	); err != nil {
		t.Fatalf("The Configuration %q has not updated LatestCreatedRevisionName from %q: %v", configName, oldRevName, err)
	}

	config, err := clients.ServingClient.Configs.Get(configName, metav1.GetOptions{})
	if err != nil {
		t.Fatal("Failed to get Configuration after it was seen to be live:", err)
	}

	return config.Status.LatestCreatedRevisionName
}

func privateServiceName(t *testing.T, clients *test.Clients, revisionName string) string {
	var privateServiceName string

	if err := wait.PollImmediate(time.Second, 1*time.Minute, func() (bool, error) {
		sks, err := clients.NetworkingClient.ServerlessServices.Get(revisionName, metav1.GetOptions{})
		if err != nil {
			return false, nil
		}
		privateServiceName = sks.Status.PrivateServiceName
		return privateServiceName != "", nil
	}); err != nil {
		t.Fatalf("Error retrieving sks %q: %v", revisionName, err)
	}

	return privateServiceName
}

// waitForDesiredScale retuns the last observed number of pods and/or error if the cond
// callback is never satisfied.
func waitForDesiredScale(clients *test.Clients, serviceName string, cond func(int) bool) (latestReady int, err error) {
	endpoints := clients.KubeClient.Kube.CoreV1().Endpoints(test.ServingNamespace)

	return latestReady, wait.PollImmediate(250*time.Millisecond, time.Minute, func() (bool, error) {
		endpoint, err := endpoints.Get(serviceName, metav1.GetOptions{})
		if err != nil {
			return false, nil
		}
		latestReady = resources.ReadyAddressCount(endpoint)
		return cond(latestReady), nil
	})
}

func ensureDesiredScale(clients *test.Clients, t *testing.T, serviceName string, cond func(int) bool) (latestReady int, observed bool) {
	endpoints := clients.KubeClient.Kube.CoreV1().Endpoints(test.ServingNamespace)

	err := wait.PollImmediate(250*time.Millisecond, 10*time.Second, func() (bool, error) {
		endpoint, err := endpoints.Get(serviceName, metav1.GetOptions{})
		if err != nil {
			return false, nil
		}

		if latestReady = resources.ReadyAddressCount(endpoint); !cond(latestReady) {
			return false, fmt.Errorf("scale %d didn't meet condition", latestReady)
		}

		return false, nil
	})
	if err != wait.ErrWaitTimeout {
		t.Log("PollError =", err)
	}

	return latestReady, err == wait.ErrWaitTimeout
}
<commit_msg>[master] Fix spelling errors (#9050)<commit_after>// +build e2e

/*
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 e2e

import (
	"fmt"
	"strconv"
	"testing"
	"time"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"knative.dev/serving/pkg/apis/autoscaling"
	v1 "knative.dev/serving/pkg/apis/serving/v1"
	"knative.dev/serving/pkg/resources"
	rtesting "knative.dev/serving/pkg/testing/v1"
	"knative.dev/serving/test"
	v1test "knative.dev/serving/test/v1"
)

func TestMinScale(t *testing.T) {
	t.Parallel()

	const minScale = 4

	clients := Setup(t)

	names := test.ResourceNames{
		// Config and Route have different names to avoid false positives
		Config: test.ObjectNameForTest(t),
		Route:  test.ObjectNameForTest(t),
		Image:  "helloworld",
	}

	test.EnsureTearDown(t, clients, &names)

	t.Log("Creating route")
	if _, err := v1test.CreateRoute(t, clients, names); err != nil {
		t.Fatal("Failed to create Route:", err)
	}

	t.Log("Creating configuration")
	cfg, err := v1test.CreateConfiguration(t, clients, names, withMinScale(minScale))
	if err != nil {
		t.Fatal("Failed to create Configuration:", err)
	}

	revName := latestRevisionName(t, clients, names.Config, "")
	serviceName := privateServiceName(t, clients, revName)

	t.Log("Waiting for revision to scale to minScale before becoming ready")
	if lr, err := waitForDesiredScale(clients, serviceName, gte(minScale)); err != nil {
		t.Fatalf("The revision %q scaled to %d < %d before becoming ready: %v", revName, lr, minScale, err)
	}

	t.Log("Waiting for revision to become ready")
	if err := v1test.WaitForRevisionState(
		clients.ServingClient, revName, v1test.IsRevisionReady, "RevisionIsReady",
	); err != nil {
		t.Fatalf("The Revision %q did not become ready: %v", revName, err)
	}

	t.Log("Holding revision at minScale after becoming ready")
	if lr, ok := ensureDesiredScale(clients, t, serviceName, gte(minScale)); !ok {
		t.Fatalf("The revision %q observed scale %d < %d after becoming ready", revName, lr, minScale)
	}

	t.Log("Updating configuration")
	if _, err := v1test.PatchConfig(t, clients, cfg, rtesting.WithConfigEnv(corev1.EnvVar{Name: "FOO", Value: "BAR"})); err != nil {
		t.Fatal("Failed to update Configuration:", err)
	}

	newRevName := latestRevisionName(t, clients, names.Config, revName)
	newServiceName := privateServiceName(t, clients, newRevName)

	t.Log("Waiting for new revision to scale to minScale after update")
	if lr, err := waitForDesiredScale(clients, newServiceName, gte(minScale)); err != nil {
		t.Fatalf("The revision %q scaled to %d < %d after creating route: %v", newRevName, lr, minScale, err)
	}

	t.Log("Waiting for new revision to become ready")
	if err := v1test.WaitForRevisionState(
		clients.ServingClient, newRevName, v1test.IsRevisionReady, "RevisionIsReady",
	); err != nil {
		t.Fatalf("The Revision %q did not become ready: %v", newRevName, err)
	}

	t.Log("Holding new revision at minScale after becoming ready")
	if lr, ok := ensureDesiredScale(clients, t, newServiceName, gte(minScale)); !ok {
		t.Fatalf("The revision %q observed scale %d < %d after becoming ready", newRevName, lr, minScale)
	}

	t.Log("Waiting for old revision to scale below minScale after being replaced")
	if lr, err := waitForDesiredScale(clients, serviceName, lt(minScale)); err != nil {
		t.Fatalf("The revision %q scaled to %d > %d after not being routable anymore: %v", revName, lr, minScale, err)
	}

	t.Log("Deleting route", names.Route)
	if err := clients.ServingClient.Routes.Delete(names.Route, &metav1.DeleteOptions{}); err != nil {
		t.Fatalf("Failed to delete route %q: %v", names.Route, err)
	}

	t.Log("Waiting for new revision to scale below minScale when there is no route")
	if lr, err := waitForDesiredScale(clients, newServiceName, lt(minScale)); err != nil {
		t.Fatalf("The revision %q scaled to %d > %d after not being routable anymore: %v", newRevName, lr, minScale, err)
	}
}

func gte(m int) func(int) bool {
	return func(n int) bool {
		return n >= m
	}
}

func lt(m int) func(int) bool {
	return func(n int) bool {
		return n < m
	}
}

func withMinScale(minScale int) func(cfg *v1.Configuration) {
	return func(cfg *v1.Configuration) {
		if cfg.Spec.Template.Annotations == nil {
			cfg.Spec.Template.Annotations = make(map[string]string, 1)
		}
		cfg.Spec.Template.Annotations[autoscaling.MinScaleAnnotationKey] = strconv.Itoa(minScale)
	}
}

func latestRevisionName(t *testing.T, clients *test.Clients, configName, oldRevName string) string {
	// Wait for the Config have a LatestCreatedRevisionName
	if err := v1test.WaitForConfigurationState(
		clients.ServingClient, configName,
		func(c *v1.Configuration) (bool, error) {
			return c.Status.LatestCreatedRevisionName != oldRevName, nil
		}, "ConfigurationHasUpdatedCreatedRevision",
	); err != nil {
		t.Fatalf("The Configuration %q has not updated LatestCreatedRevisionName from %q: %v", configName, oldRevName, err)
	}

	config, err := clients.ServingClient.Configs.Get(configName, metav1.GetOptions{})
	if err != nil {
		t.Fatal("Failed to get Configuration after it was seen to be live:", err)
	}

	return config.Status.LatestCreatedRevisionName
}

func privateServiceName(t *testing.T, clients *test.Clients, revisionName string) string {
	var privateServiceName string

	if err := wait.PollImmediate(time.Second, 1*time.Minute, func() (bool, error) {
		sks, err := clients.NetworkingClient.ServerlessServices.Get(revisionName, metav1.GetOptions{})
		if err != nil {
			return false, nil
		}
		privateServiceName = sks.Status.PrivateServiceName
		return privateServiceName != "", nil
	}); err != nil {
		t.Fatalf("Error retrieving sks %q: %v", revisionName, err)
	}

	return privateServiceName
}

// waitForDesiredScale returns the last observed number of pods and/or error if the cond
// callback is never satisfied.
func waitForDesiredScale(clients *test.Clients, serviceName string, cond func(int) bool) (latestReady int, err error) {
	endpoints := clients.KubeClient.Kube.CoreV1().Endpoints(test.ServingNamespace)

	return latestReady, wait.PollImmediate(250*time.Millisecond, time.Minute, func() (bool, error) {
		endpoint, err := endpoints.Get(serviceName, metav1.GetOptions{})
		if err != nil {
			return false, nil
		}
		latestReady = resources.ReadyAddressCount(endpoint)
		return cond(latestReady), nil
	})
}

func ensureDesiredScale(clients *test.Clients, t *testing.T, serviceName string, cond func(int) bool) (latestReady int, observed bool) {
	endpoints := clients.KubeClient.Kube.CoreV1().Endpoints(test.ServingNamespace)

	err := wait.PollImmediate(250*time.Millisecond, 10*time.Second, func() (bool, error) {
		endpoint, err := endpoints.Get(serviceName, metav1.GetOptions{})
		if err != nil {
			return false, nil
		}

		if latestReady = resources.ReadyAddressCount(endpoint); !cond(latestReady) {
			return false, fmt.Errorf("scale %d didn't meet condition", latestReady)
		}

		return false, nil
	})
	if err != wait.ErrWaitTimeout {
		t.Log("PollError =", err)
	}

	return latestReady, err == wait.ErrWaitTimeout
}
<|endoftext|>


<commit_before>package reflection

import (
	"bytes"
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"reflect"
	"sync"

	"github.com/golang/protobuf/proto"
	dpb "github.com/golang/protobuf/protoc-gen-go/descriptor"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	rpb "google.golang.org/grpc/reflection/grpc_reflection_v1alpha"
)

type serverReflectionServer struct {
	s *grpc.Server
	// TODO mu is not used. Add lock() and unlock().
	mu                sync.Mutex
	typeToNameMap     map[reflect.Type]string
	nameToTypeMap     map[string]reflect.Type
	typeToFileDescMap map[reflect.Type]*dpb.FileDescriptorProto
	// TODO remove this, replace with s.ftdmap
	filenameToDescMap map[string]*dpb.FileDescriptorProto
}

// InstallOnServer installs server reflection service on the given grpc server.
func InstallOnServer(s *grpc.Server) {
	rpb.RegisterServerReflectionServer(s, &serverReflectionServer{
		s:                 s,
		typeToNameMap:     make(map[reflect.Type]string),
		nameToTypeMap:     make(map[string]reflect.Type),
		typeToFileDescMap: make(map[reflect.Type]*dpb.FileDescriptorProto),
		filenameToDescMap: make(map[string]*dpb.FileDescriptorProto),
	})
}

type protoMessage interface {
	Descriptor() ([]byte, []int)
}

func (s *serverReflectionServer) fileDescForType(st reflect.Type) (*dpb.FileDescriptorProto, []int, error) {
	// Indexes list is not stored in cache.
	// So this step is needed to get idxs.
	m, ok := reflect.Zero(reflect.PtrTo(st)).Interface().(protoMessage)
	if !ok {
		return nil, nil, fmt.Errorf("failed to create message from type: %v", st)
	}
	enc, idxs := m.Descriptor()

	// Check type to fileDesc cache.
	if fd, ok := s.typeToFileDescMap[st]; ok {
		return fd, idxs, nil
	}

	// Cache missed, try to decode.
	fd, err := s.decodeFileDesc(enc)
	if err != nil {
		return nil, nil, err
	}
	// Add to cache.
	s.typeToFileDescMap[st] = fd
	return fd, idxs, nil
}

func (s *serverReflectionServer) decodeFileDesc(enc []byte) (*dpb.FileDescriptorProto, error) {
	raw := decompress(enc)
	if raw == nil {
		return nil, fmt.Errorf("failed to decompress enc")
	}

	fd := new(dpb.FileDescriptorProto)
	if err := proto.Unmarshal(raw, fd); err != nil {
		return nil, fmt.Errorf("bad descriptor: %v", err)
	}
	// If decodeFileDesc is called, it's the first time this file is seen.
	// Add it to cache.
	s.filenameToDescMap[fd.GetName()] = fd
	return fd, nil
}

func decompress(b []byte) []byte {
	r, err := gzip.NewReader(bytes.NewReader(b))
	if err != nil {
		fmt.Printf("bad gzipped descriptor: %v\n", err)
		return nil
	}
	out, err := ioutil.ReadAll(r)
	if err != nil {
		fmt.Printf("bad gzipped descriptor: %v\n", err)
		return nil
	}
	return out
}

func (s *serverReflectionServer) typeForName(name string) (reflect.Type, error) {
	// Check cache first.
	if st, ok := s.nameToTypeMap[name]; ok {
		return st, nil
	}

	pt := proto.MessageType(name)
	if pt == nil {
		return nil, fmt.Errorf("unknown type: %q", name)
	}
	st := pt.Elem()

	// Add to cache.
	s.typeToNameMap[st] = name
	s.nameToTypeMap[name] = st

	// TODO is this necessary?
	// In most cases, the returned type will be used to search
	// for file descriptor.
	// Add it to cache now.
	fd, _, err := s.fileDescForType(st)
	if err == nil {
		s.typeToFileDescMap[st] = fd
	}

	return st, nil
}

func (s *serverReflectionServer) nameForType(st reflect.Type) (string, error) {
	// Check cache first.
	if name, ok := s.typeToNameMap[st]; ok {
		return name, nil
	}

	var name string
	fd, idxs, err := s.fileDescForType(st)
	if err != nil {
		return "", err
	}
	mt := fd.MessageType[idxs[0]]
	name = mt.GetName()
	for i := 1; i < len(idxs); i++ {
		mt = mt.NestedType[idxs[i]]
		name += "_" + mt.GetName()
	}
	if fd.Package != nil {
		name = *fd.Package + "." + name
	}

	// Add to cache.
	s.typeToNameMap[st] = name
	s.nameToTypeMap[name] = st

	return name, nil
}

func (s *serverReflectionServer) nameForPointer(i interface{}) (string, error) {
	return s.nameForType(reflect.TypeOf(i).Elem())
}

func (s *serverReflectionServer) filenameForType(st reflect.Type) (string, error) {
	// Check cache first. The value of cache is descriptor, not filename.
	if fd, ok := s.typeToFileDescMap[st]; ok {
		return fd.GetName(), nil
	}

	fd, _, err := s.fileDescForType(st)
	if err != nil {
		return "", err
	}
	return fd.GetName(), nil
}

// TODO filenameForMethod and Service

func (s *serverReflectionServer) fileDescContainingExtension(st reflect.Type, ext int32) (*dpb.FileDescriptorProto, error) {
	m, ok := reflect.Zero(reflect.PtrTo(st)).Interface().(proto.Message)
	if !ok {
		return nil, fmt.Errorf("failed to create message from type: %v", st)
	}

	var extDesc *proto.ExtensionDesc
	for id, desc := range proto.RegisteredExtensions(m) {
		if id == ext {
			extDesc = desc
			break
		}
	}

	if extDesc == nil {
		return nil, fmt.Errorf("failed to find registered extension for extension number %v", ext)
	}

	extT := reflect.TypeOf(extDesc.ExtensionType).Elem()
	// TODO this doesn't work if extT is simple types, like int32
	// Check cache.
	if fd, ok := s.typeToFileDescMap[extT]; ok {
		return fd, nil
	}

	fd, _, err := s.fileDescForType(extT)
	if err != nil {
		return nil, err
	}
	return fd, nil
}

// TODO filenameContainingExtension
// fd := fileDescContainingExtension()
// return fd.GetName()

// fileDescWireFormatByFilename returns the file descriptor of file with the given name.
// TODO exporte and add lock
func (s *serverReflectionServer) fileDescWireFormatByFilename(name string) ([]byte, error) {
	fd, ok := s.filenameToDescMap[name]
	if !ok {
		return nil, fmt.Errorf("unknown file: %v", name)
	}
	b, err := proto.Marshal(fd)
	if err != nil {
		return nil, err
	}
	return b, nil
}

func (s *serverReflectionServer) allExtensionNumbersForType(st reflect.Type) ([]int32, error) {
	m, ok := reflect.Zero(reflect.PtrTo(st)).Interface().(proto.Message)
	if !ok {
		return nil, fmt.Errorf("failed to create message from type: %v", st)
	}

	var out []int32
	for id := range proto.RegisteredExtensions(m) {
		out = append(out, id)
	}
	return out, nil
}

func (s *serverReflectionServer) ServerReflectionInfo(stream rpb.ServerReflection_ServerReflectionInfoServer) error {
	for {
		in, err := stream.Recv()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}

		var response *rpb.FileDescriptorResponse
		switch req := in.MessageRequest.(type) {
		case *rpb.ServerReflectionRequest_FileByFilename:
			b, err := s.fileDescWireFormatByFilename(req.FileByFilename)
			if err != nil {
				// TODO grpc error or send message back
				return err
			}
			response = &rpb.FileDescriptorResponse{FileDescriptorProto: [][]byte{b}}
		case *rpb.ServerReflectionRequest_FileContainingSymbol:
		case *rpb.ServerReflectionRequest_FileContainingExtension:
		case *rpb.ServerReflectionRequest_AllExtensionNumbersOfType:
		case *rpb.ServerReflectionRequest_ListServices:
		default:
			return grpc.Errorf(codes.InvalidArgument, "invalid MessageRequest: %v", in.MessageRequest)
		}

		out := &rpb.ServerReflectionResponse{
			ValidHost:       in.Host,
			OriginalRequest: in,
			MessageResponse: &rpb.ServerReflectionResponse_FileDescriptorResponse{
				FileDescriptorResponse: response,
			},
		}
		if err := stream.Send(out); err != nil {
			return err
		}
	}
}
<commit_msg>Add fileDescWireFormatContainingSymbol and handle file_containing_symbol<commit_after>package reflection

import (
	"bytes"
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"reflect"
	"sync"

	"github.com/golang/protobuf/proto"
	dpb "github.com/golang/protobuf/protoc-gen-go/descriptor"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	rpb "google.golang.org/grpc/reflection/grpc_reflection_v1alpha"
)

type serverReflectionServer struct {
	s *grpc.Server
	// TODO mu is not used. Add lock() and unlock().
	mu                sync.Mutex
	typeToNameMap     map[reflect.Type]string
	nameToTypeMap     map[string]reflect.Type
	typeToFileDescMap map[reflect.Type]*dpb.FileDescriptorProto
	// TODO remove this, replace with s.ftdmap
	filenameToDescMap map[string]*dpb.FileDescriptorProto
}

// InstallOnServer installs server reflection service on the given grpc server.
func InstallOnServer(s *grpc.Server) {
	rpb.RegisterServerReflectionServer(s, &serverReflectionServer{
		s:                 s,
		typeToNameMap:     make(map[reflect.Type]string),
		nameToTypeMap:     make(map[string]reflect.Type),
		typeToFileDescMap: make(map[reflect.Type]*dpb.FileDescriptorProto),
		filenameToDescMap: make(map[string]*dpb.FileDescriptorProto),
	})
}

type protoMessage interface {
	Descriptor() ([]byte, []int)
}

func (s *serverReflectionServer) fileDescForType(st reflect.Type) (*dpb.FileDescriptorProto, []int, error) {
	// Indexes list is not stored in cache.
	// So this step is needed to get idxs.
	m, ok := reflect.Zero(reflect.PtrTo(st)).Interface().(protoMessage)
	if !ok {
		return nil, nil, fmt.Errorf("failed to create message from type: %v", st)
	}
	enc, idxs := m.Descriptor()

	// Check type to fileDesc cache.
	if fd, ok := s.typeToFileDescMap[st]; ok {
		return fd, idxs, nil
	}

	// Cache missed, try to decode.
	fd, err := s.decodeFileDesc(enc)
	if err != nil {
		return nil, nil, err
	}
	// Add to cache.
	s.typeToFileDescMap[st] = fd
	return fd, idxs, nil
}

func (s *serverReflectionServer) decodeFileDesc(enc []byte) (*dpb.FileDescriptorProto, error) {
	raw := decompress(enc)
	if raw == nil {
		return nil, fmt.Errorf("failed to decompress enc")
	}

	fd := new(dpb.FileDescriptorProto)
	if err := proto.Unmarshal(raw, fd); err != nil {
		return nil, fmt.Errorf("bad descriptor: %v", err)
	}
	// If decodeFileDesc is called, it's the first time this file is seen.
	// Add it to cache.
	s.filenameToDescMap[fd.GetName()] = fd
	return fd, nil
}

func decompress(b []byte) []byte {
	r, err := gzip.NewReader(bytes.NewReader(b))
	if err != nil {
		fmt.Printf("bad gzipped descriptor: %v\n", err)
		return nil
	}
	out, err := ioutil.ReadAll(r)
	if err != nil {
		fmt.Printf("bad gzipped descriptor: %v\n", err)
		return nil
	}
	return out
}

func (s *serverReflectionServer) typeForName(name string) (reflect.Type, error) {
	// Check cache first.
	if st, ok := s.nameToTypeMap[name]; ok {
		return st, nil
	}

	pt := proto.MessageType(name)
	if pt == nil {
		return nil, fmt.Errorf("unknown type: %q", name)
	}
	st := pt.Elem()

	// Add to cache.
	s.typeToNameMap[st] = name
	s.nameToTypeMap[name] = st

	// TODO is this necessary?
	// In most cases, the returned type will be used to search
	// for file descriptor.
	// Add it to cache now.
	fd, _, err := s.fileDescForType(st)
	if err == nil {
		s.typeToFileDescMap[st] = fd
	}

	return st, nil
}

func (s *serverReflectionServer) nameForType(st reflect.Type) (string, error) {
	// Check cache first.
	if name, ok := s.typeToNameMap[st]; ok {
		return name, nil
	}

	var name string
	fd, idxs, err := s.fileDescForType(st)
	if err != nil {
		return "", err
	}
	mt := fd.MessageType[idxs[0]]
	name = mt.GetName()
	for i := 1; i < len(idxs); i++ {
		mt = mt.NestedType[idxs[i]]
		name += "_" + mt.GetName()
	}
	if fd.Package != nil {
		name = *fd.Package + "." + name
	}

	// Add to cache.
	s.typeToNameMap[st] = name
	s.nameToTypeMap[name] = st

	return name, nil
}

func (s *serverReflectionServer) nameForPointer(i interface{}) (string, error) {
	return s.nameForType(reflect.TypeOf(i).Elem())
}

func (s *serverReflectionServer) filenameForType(st reflect.Type) (string, error) {
	// Check cache first. The value of cache is descriptor, not filename.
	if fd, ok := s.typeToFileDescMap[st]; ok {
		return fd.GetName(), nil
	}

	fd, _, err := s.fileDescForType(st)
	if err != nil {
		return "", err
	}
	return fd.GetName(), nil
}

// TODO filenameForMethod and Service

func (s *serverReflectionServer) fileDescContainingExtension(st reflect.Type, ext int32) (*dpb.FileDescriptorProto, error) {
	m, ok := reflect.Zero(reflect.PtrTo(st)).Interface().(proto.Message)
	if !ok {
		return nil, fmt.Errorf("failed to create message from type: %v", st)
	}

	var extDesc *proto.ExtensionDesc
	for id, desc := range proto.RegisteredExtensions(m) {
		if id == ext {
			extDesc = desc
			break
		}
	}

	if extDesc == nil {
		return nil, fmt.Errorf("failed to find registered extension for extension number %v", ext)
	}

	extT := reflect.TypeOf(extDesc.ExtensionType).Elem()
	// TODO this doesn't work if extT is simple types, like int32
	// Check cache.
	if fd, ok := s.typeToFileDescMap[extT]; ok {
		return fd, nil
	}

	fd, _, err := s.fileDescForType(extT)
	if err != nil {
		return nil, err
	}
	return fd, nil
}

// TODO filenameContainingExtension
// fd := fileDescContainingExtension()
// return fd.GetName()

// fileDescWireFormatByFilename returns the file descriptor of file with the given name.
// TODO exporte and add lock
func (s *serverReflectionServer) fileDescWireFormatByFilename(name string) ([]byte, error) {
	fd, ok := s.filenameToDescMap[name]
	if !ok {
		return nil, fmt.Errorf("unknown file: %v", name)
	}
	b, err := proto.Marshal(fd)
	if err != nil {
		return nil, err
	}
	return b, nil
}

func (s *serverReflectionServer) fileDescWireFormatContainingSymbol(name string) ([]byte, error) {
	var (
		fd *dpb.FileDescriptorProto
	)
	// Check if it's a type name.
	if st, err := s.typeForName(name); err == nil {
		fd, _, err = s.fileDescForType(st)
		if err != nil {
			return nil, err
		}
	} else {
		// Check if it's a service name or method name.
		meta := s.s.Metadata(name)
		if meta != nil {
			if enc, ok := meta.([]byte); ok {
				fd, err = s.decodeFileDesc(enc)
				if err != nil {
					return nil, err
				}
			}
		}
	}

	// Marshal to wire format.
	if fd != nil {
		b, err := proto.Marshal(fd)
		if err != nil {
			return nil, err
		}
		return b, nil
	}
	return nil, fmt.Errorf("unknown symbol: %v", name)
}

func (s *serverReflectionServer) allExtensionNumbersForType(st reflect.Type) ([]int32, error) {
	m, ok := reflect.Zero(reflect.PtrTo(st)).Interface().(proto.Message)
	if !ok {
		return nil, fmt.Errorf("failed to create message from type: %v", st)
	}

	var out []int32
	for id := range proto.RegisteredExtensions(m) {
		out = append(out, id)
	}
	return out, nil
}

func (s *serverReflectionServer) ServerReflectionInfo(stream rpb.ServerReflection_ServerReflectionInfoServer) error {
	for {
		in, err := stream.Recv()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}

		var response *rpb.FileDescriptorResponse
		switch req := in.MessageRequest.(type) {
		case *rpb.ServerReflectionRequest_FileByFilename:
			b, err := s.fileDescWireFormatByFilename(req.FileByFilename)
			if err != nil {
				// TODO grpc error or send message back
				return err
			}
			response = &rpb.FileDescriptorResponse{FileDescriptorProto: [][]byte{b}}
		case *rpb.ServerReflectionRequest_FileContainingSymbol:
			b, err := s.fileDescWireFormatContainingSymbol(req.FileContainingSymbol)
			if err != nil {
				// TODO grpc error or send message back
				return err
			}
			response = &rpb.FileDescriptorResponse{FileDescriptorProto: [][]byte{b}}
		case *rpb.ServerReflectionRequest_FileContainingExtension:
		case *rpb.ServerReflectionRequest_AllExtensionNumbersOfType:
		case *rpb.ServerReflectionRequest_ListServices:
		default:
			return grpc.Errorf(codes.InvalidArgument, "invalid MessageRequest: %v", in.MessageRequest)
		}

		out := &rpb.ServerReflectionResponse{
			ValidHost:       in.Host,
			OriginalRequest: in,
			MessageResponse: &rpb.ServerReflectionResponse_FileDescriptorResponse{
				FileDescriptorResponse: response,
			},
		}
		if err := stream.Send(out); err != nil {
			return err
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"regexp"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ecs"
	"fmt"
	"os"
	"time"
)

func ExtractName(a *string) string {
	re := regexp.MustCompile("\\w+:\\w+:\\w+:[\\w-]+:\\d+:[^/]+\\/(.+)")

	res := re.FindAllStringSubmatch(*a, -1);

	if res == nil {
		return "No match for " + *a;
	}

	return res[0][1]
}

func ExtractVersion(a string) (string, string) {
	re := regexp.MustCompile("(.*?)\\:(.+)")

	res := re.FindAllStringSubmatch(a, -1);

	if (res == nil) {
		return "", ""
	}

	return res[0][2], res[0][1]
}

func _listServices(cluster string) *ecs.ListServicesOutput {
	svc := ecs.New(session.New(), _getAwsConfig())

	params := &ecs.ListServicesInput{
		Cluster: aws.String(cluster),
		MaxResults: aws.Int64(10),
	}

	resp, err := svc.ListServices(params)
	if err != nil {
		// Print the error, cast err to awserr.Error to get the Code and
		// Message from an error.
		fmt.Println(err.Error())
		os.Exit(1);
	}

	return resp;
}

func _listTasks(cluster, service string) *ecs.ListTasksOutput {
	svc := ecs.New(session.New(), _getAwsConfig())

	params := &ecs.ListTasksInput{
		Cluster: aws.String(cluster),
		ServiceName: aws.String(service),
	}

	resp, err := svc.ListTasks(params)
	if (err != nil) {
		fmt.Println(err.Error())
		os.Exit(1)
	}

	return resp
}

func _describeService(clusterArn, serviceArn string) (*ecs.DescribeServicesOutput) {
	svc := ecs.New(session.New(), _getAwsConfig())

	params := &ecs.DescribeServicesInput{
		Cluster: aws.String(clusterArn),
		Services: []*string{aws.String(serviceArn)},
	}

	result, err := svc.DescribeServices(params)
	if (err != nil) {
		errState(err.Error())
	}

	return result;
}

func _describeTasks(clusterArn string, tasks []*string) *ecs.DescribeTasksOutput {
	svc := ecs.New(session.New(), _getAwsConfig())

	params := &ecs.DescribeTasksInput{
		Cluster: aws.String(clusterArn),
		Tasks: tasks,
	}

	result, err := svc.DescribeTasks(params)

	if (err != nil) {
		errState(err.Error())
	}

	return result
}

func _createTaskDefinition(taskDefinition *ecs.DescribeTaskDefinitionOutput) string {
	svc := ecs.New(session.New(), _getAwsConfig())

	params := &ecs.RegisterTaskDefinitionInput{
		ContainerDefinitions : taskDefinition.TaskDefinition.ContainerDefinitions,
		Family: taskDefinition.TaskDefinition.Family,
		Volumes: taskDefinition.TaskDefinition.Volumes,
	}

	registrationResult, err := svc.RegisterTaskDefinition(params)

	if (err != nil) {
		errState(err.Error())
	}

	taskDefinitionArn := registrationResult.TaskDefinition.TaskDefinitionArn;

	return *taskDefinitionArn
}

func _updateService(newTaskDefinitionArn string, service *ecs.DescribeServicesOutput) {
	svc := ecs.New(session.New(), _getAwsConfig())

	params := &ecs.UpdateServiceInput{
		Cluster: aws.String(*service.Services[0].ClusterArn),
		DesiredCount: aws.Int64(*service.Services[0].DesiredCount),
		Service: aws.String(*service.Services[0].ServiceArn),
		TaskDefinition: aws.String(newTaskDefinitionArn),
	}

	_, err := svc.UpdateService(params)

	if (err != nil) {
		errState(err.Error())
	}
}

func ListServices(clusterArn string) {

	resp := _listServices(clusterArn)
	for i := 0; i < len(resp.ServiceArns); i++ {
		name := ExtractName(resp.ServiceArns[i])
		fmt.Println(name)
	}
}

func ListTasks(clusterArn, serviceArn string) {
	resp := _listTasks(clusterArn, serviceArn)
	for i := 0; i < len(resp.TaskArns); i++ {
		name := ExtractName(resp.TaskArns[i])
		fmt.Println(name)
	}
}

func UpdateService(clusterArn, serviceArn string) {
	tasks := _listTasks(clusterArn, serviceArn)
	service := _describeService(clusterArn, serviceArn)

	if (len(service.Services) > 1) {
		errState("No support for multiple services.");
	}

	desiredCount := *service.Services[0].DesiredCount

	if len(tasks.TaskArns) < int(desiredCount) {
		errStatef("The number of actual tasks (%d) differs from desired tasks (%d)\n", len(tasks.TaskArns), desiredCount)
	}

	for i := 0; i < len(tasks.TaskArns); i++ {
		currentTasks := _listTasks(clusterArn, serviceArn)
		fmt.Println("Task: " + *tasks.TaskArns[i])
		_stopTask(clusterArn, *tasks.TaskArns[i])
		fmt.Println("  Stopped")
		_waitForNewTask(clusterArn, serviceArn, currentTasks.TaskArns)
	}

	fmt.Printf("Service [%s] is updated\n", *service.Services[0].ServiceName)
}

func DescribeService(clusterArn, serviceArn string) {

	service := _describeService(clusterArn, serviceArn)

	for n := 0; n < len(service.Services); n++ {
		item := service.Services[n]

		fmt.Printf("Service name [%s], Running: %d, Pending: %d, Desired: %d\n", *item.ServiceName, *item.RunningCount, *item.PendingCount, *item.DesiredCount)

		if (verboseLevel == 1) {
			for i := 0; i < len(item.Deployments); i++ {
				deployment := item.Deployments[i]
				fmt.Printf("   %s (%s), running: %d, Pending: %d: Desired: %d\n", ExtractName(deployment.TaskDefinition), *deployment.Status, *deployment.RunningCount, *deployment.PendingCount, *deployment.DesiredCount)
			}
		}

		if (verboseLevel == 2) {
			fmt.Printf("Deployment configuration -> MaximumPercent: %d, MinimumHealthyPercent: %d\n", *item.DeploymentConfiguration.MaximumPercent, *item.DeploymentConfiguration.MinimumHealthyPercent)
			definition := _describeTaskDefinition(*item.TaskDefinition)
			fmt.Println(definition);
		}
	}
}

func ReleaseService(clusterArn, serviceArn, version string) {
	service := _describeService(clusterArn, serviceArn);

	if (len(service.Services) > 1) {
		errState("No support for multiple services.");
	}

	taskDefinitionName := *service.Services[0].TaskDefinition
	taskDefinition := _describeTaskDefinition(taskDefinitionName)
	dockerImage := *taskDefinition.TaskDefinition.ContainerDefinitions[0].Image
	currentVersion, imagePart := ExtractVersion(dockerImage)

	fmt.Printf("Service            [%s]\n", *service.Services[0].ServiceName);
	fmt.Printf("Task definition    [%s]\n", taskDefinitionName);
	fmt.Printf("Docker image       [%s]\n", imagePart)
	fmt.Printf("Version            [%s]\n", currentVersion);

	if (version == currentVersion) {
		errUsage("Specified version is already deployed!");
	}

	minimumHealthyPercentage := *service.Services[0].DeploymentConfiguration.MinimumHealthyPercent
	desiredCount := *service.Services[0].DesiredCount
	minimumHealthyCount := desiredCount * minimumHealthyPercentage / 100

	fmt.Printf("HealthyPercentage [%d], DesiredCount [%d] -> Number of running instances during deployment [%d] ... ", minimumHealthyPercentage, desiredCount, minimumHealthyCount)

	if (desiredCount - minimumHealthyCount == 0) {
		errUsage("Not possible to deploy because of too high healthy percentage")
	}

	fmt.Println("OK")

	*taskDefinition.TaskDefinition.ContainerDefinitions[0].Image = imagePart + ":" + version

	fmt.Println(*taskDefinition.TaskDefinition.ContainerDefinitions[0].Image)

	newTaskDefinitionArn := _createTaskDefinition(taskDefinition)

	_updateService(newTaskDefinitionArn, service)
}

func _describeTaskDefinition(taskDefinitionName string) *ecs.DescribeTaskDefinitionOutput {
	svc := ecs.New(session.New(), _getAwsConfig())

	params := &ecs.DescribeTaskDefinitionInput{
		TaskDefinition: aws.String(taskDefinitionName),
	}

	result, err := svc.DescribeTaskDefinition(params)

	if (err != nil) {
		errState(err.Error())
	}

	return result
}

func _stopTask(cluster, taskArn string) {
	svc := ecs.New(session.New(), _getAwsConfig())

	params := &ecs.StopTaskInput{
		Cluster: aws.String(cluster),
		Reason: aws.String("Stopped by WriterTool"),
		Task: aws.String(taskArn),
	}

	_, err := svc.StopTask(params)

	if (err != nil) {
		errState(err.Error())
	}
}

func _waitForNewTask(cluster string, service string, tasks []*string) {
	fmt.Print("  Waiting for new task ")
	newTask := ""

	attempts := 60
	sleepTime := 2

	for i := 0; i < attempts; i++ {
		currentTasks := _listTasks(cluster, service)
		if (len(currentTasks.TaskArns) < len(tasks) || len(currentTasks.TaskArns) == 0) {
			fmt.Print(".")
		} else if (newTask == "") {
			newTask = _findNewTask(tasks, currentTasks.TaskArns)
			if (newTask == "") {
				errState("No new task found among tasks")
			}
			fmt.Println(" done")
			fmt.Print("  New task: " + newTask + " ")

		}

		if (newTask != "") {
			taskStates := _describeTasks(cluster, []*string{aws.String(newTask)})
			if (*taskStates.Tasks[0].LastStatus == "RUNNING") {
				fmt.Println(" done")
				return
			} else if (*taskStates.Tasks[0].LastStatus == "STOPPED") {
				fmt.Print("X\n")
				errState("Could not start new task")
			} else {
				fmt.Print(".")
			}
		}

		time.Sleep(time.Duration(sleepTime) * time.Second)
	}

	errStatef("Task [%s] did not start in %d attempts", newTask, attempts)

}

func _findNewTask(tasks, currentTasks []*string) string {
	for i := 0; i < len(tasks); i++ {
		for j := 0; j < len(currentTasks); j++ {
			if *tasks[i] != *currentTasks[j] {
				return *currentTasks[j]
			}
		}
	}

	return ""
}

func GetServiceArn(clusterArn, name string) string {
	clusterArns := _listServices(clusterArn);

	for i := 0; i < len(clusterArns.ServiceArns); i++ {
		arn := clusterArns.ServiceArns[i];
		if (ExtractName(arn) == name) {
			return *arn;
		}
	}

	return "";
}
<commit_msg>Fixing issue where no new tasks where found during service update<commit_after>package main

import (
	"regexp"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ecs"
	"fmt"
	"os"
	"time"
)

func ExtractName(a *string) string {
	re := regexp.MustCompile("\\w+:\\w+:\\w+:[\\w-]+:\\d+:[^/]+\\/(.+)")

	res := re.FindAllStringSubmatch(*a, -1);

	if res == nil {
		return "No match for " + *a;
	}

	return res[0][1]
}

func ExtractVersion(a string) (string, string) {
	re := regexp.MustCompile("(.*?)\\:(.+)")

	res := re.FindAllStringSubmatch(a, -1);

	if (res == nil) {
		return "", ""
	}

	return res[0][2], res[0][1]
}

func _listServices(cluster string) *ecs.ListServicesOutput {
	svc := ecs.New(session.New(), _getAwsConfig())

	params := &ecs.ListServicesInput{
		Cluster: aws.String(cluster),
		MaxResults: aws.Int64(10),
	}

	resp, err := svc.ListServices(params)
	if err != nil {
		// Print the error, cast err to awserr.Error to get the Code and
		// Message from an error.
		fmt.Println(err.Error())
		os.Exit(1);
	}

	return resp;
}

func _listTasks(cluster, service string) *ecs.ListTasksOutput {
	svc := ecs.New(session.New(), _getAwsConfig())

	params := &ecs.ListTasksInput{
		Cluster: aws.String(cluster),
		ServiceName: aws.String(service),
	}

	resp, err := svc.ListTasks(params)
	if (err != nil) {
		fmt.Println(err.Error())
		os.Exit(1)
	}

	return resp
}

func _describeService(clusterArn, serviceArn string) (*ecs.DescribeServicesOutput) {
	svc := ecs.New(session.New(), _getAwsConfig())

	params := &ecs.DescribeServicesInput{
		Cluster: aws.String(clusterArn),
		Services: []*string{aws.String(serviceArn)},
	}

	result, err := svc.DescribeServices(params)
	if (err != nil) {
		errState(err.Error())
	}

	return result;
}

func _describeTasks(clusterArn string, tasks []*string) *ecs.DescribeTasksOutput {
	svc := ecs.New(session.New(), _getAwsConfig())

	params := &ecs.DescribeTasksInput{
		Cluster: aws.String(clusterArn),
		Tasks: tasks,
	}

	result, err := svc.DescribeTasks(params)

	if (err != nil) {
		errState(err.Error())
	}

	return result
}

func _createTaskDefinition(taskDefinition *ecs.DescribeTaskDefinitionOutput) string {
	svc := ecs.New(session.New(), _getAwsConfig())

	params := &ecs.RegisterTaskDefinitionInput{
		ContainerDefinitions : taskDefinition.TaskDefinition.ContainerDefinitions,
		Family: taskDefinition.TaskDefinition.Family,
		Volumes: taskDefinition.TaskDefinition.Volumes,
	}

	registrationResult, err := svc.RegisterTaskDefinition(params)

	if (err != nil) {
		errState(err.Error())
	}

	taskDefinitionArn := registrationResult.TaskDefinition.TaskDefinitionArn;

	return *taskDefinitionArn
}

func _updateService(newTaskDefinitionArn string, service *ecs.DescribeServicesOutput) {
	svc := ecs.New(session.New(), _getAwsConfig())

	params := &ecs.UpdateServiceInput{
		Cluster: aws.String(*service.Services[0].ClusterArn),
		DesiredCount: aws.Int64(*service.Services[0].DesiredCount),
		Service: aws.String(*service.Services[0].ServiceArn),
		TaskDefinition: aws.String(newTaskDefinitionArn),
	}

	_, err := svc.UpdateService(params)

	if (err != nil) {
		errState(err.Error())
	}
}

func ListServices(clusterArn string) {

	resp := _listServices(clusterArn)
	for i := 0; i < len(resp.ServiceArns); i++ {
		name := ExtractName(resp.ServiceArns[i])
		fmt.Println(name)
	}
}

func ListTasks(clusterArn, serviceArn string) {
	resp := _listTasks(clusterArn, serviceArn)
	for i := 0; i < len(resp.TaskArns); i++ {
		name := ExtractName(resp.TaskArns[i])
		fmt.Println(name)
	}
}

func UpdateService(clusterArn, serviceArn string) {
	tasks := _listTasks(clusterArn, serviceArn)
	service := _describeService(clusterArn, serviceArn)

	if (len(service.Services) > 1) {
		errState("No support for multiple services.");
	}

	desiredCount := *service.Services[0].DesiredCount

	if len(tasks.TaskArns) < int(desiredCount) {
		errStatef("The number of actual tasks (%d) differs from desired tasks (%d)\n", len(tasks.TaskArns), desiredCount)
	}

	for i := 0; i < len(tasks.TaskArns); i++ {
		currentTasks := _listTasks(clusterArn, serviceArn)
		fmt.Println("Task: " + *tasks.TaskArns[i])
		_stopTask(clusterArn, *tasks.TaskArns[i])
		fmt.Println("  Stopped")
		_waitForNewTask(clusterArn, serviceArn, currentTasks.TaskArns)
	}

	fmt.Printf("Service [%s] is updated\n", *service.Services[0].ServiceName)
}

func DescribeService(clusterArn, serviceArn string) {

	service := _describeService(clusterArn, serviceArn)

	for n := 0; n < len(service.Services); n++ {
		item := service.Services[n]

		fmt.Printf("Service name [%s], Running: %d, Pending: %d, Desired: %d\n", *item.ServiceName, *item.RunningCount, *item.PendingCount, *item.DesiredCount)

		if (verboseLevel == 1) {
			for i := 0; i < len(item.Deployments); i++ {
				deployment := item.Deployments[i]
				fmt.Printf("   %s (%s), running: %d, Pending: %d: Desired: %d\n", ExtractName(deployment.TaskDefinition), *deployment.Status, *deployment.RunningCount, *deployment.PendingCount, *deployment.DesiredCount)
			}
		}

		if (verboseLevel == 2) {
			fmt.Printf("Deployment configuration -> MaximumPercent: %d, MinimumHealthyPercent: %d\n", *item.DeploymentConfiguration.MaximumPercent, *item.DeploymentConfiguration.MinimumHealthyPercent)
			definition := _describeTaskDefinition(*item.TaskDefinition)
			fmt.Println(definition);
		}
	}
}

func ReleaseService(clusterArn, serviceArn, version string) {
	service := _describeService(clusterArn, serviceArn);

	if (len(service.Services) > 1) {
		errState("No support for multiple services.");
	}

	taskDefinitionName := *service.Services[0].TaskDefinition
	taskDefinition := _describeTaskDefinition(taskDefinitionName)
	dockerImage := *taskDefinition.TaskDefinition.ContainerDefinitions[0].Image
	currentVersion, imagePart := ExtractVersion(dockerImage)

	fmt.Printf("Service            [%s]\n", *service.Services[0].ServiceName);
	fmt.Printf("Task definition    [%s]\n", taskDefinitionName);
	fmt.Printf("Docker image       [%s]\n", imagePart)
	fmt.Printf("Version            [%s]\n", currentVersion);

	if (version == currentVersion) {
		errUsage("Specified version is already deployed!");
	}

	minimumHealthyPercentage := *service.Services[0].DeploymentConfiguration.MinimumHealthyPercent
	desiredCount := *service.Services[0].DesiredCount
	minimumHealthyCount := desiredCount * minimumHealthyPercentage / 100

	fmt.Printf("HealthyPercentage [%d], DesiredCount [%d] -> Number of running instances during deployment [%d] ... ", minimumHealthyPercentage, desiredCount, minimumHealthyCount)

	if (desiredCount - minimumHealthyCount == 0) {
		errUsage("Not possible to deploy because of too high healthy percentage")
	}

	fmt.Println("OK")

	*taskDefinition.TaskDefinition.ContainerDefinitions[0].Image = imagePart + ":" + version

	fmt.Println(*taskDefinition.TaskDefinition.ContainerDefinitions[0].Image)

	newTaskDefinitionArn := _createTaskDefinition(taskDefinition)

	_updateService(newTaskDefinitionArn, service)
}

func _describeTaskDefinition(taskDefinitionName string) *ecs.DescribeTaskDefinitionOutput {
	svc := ecs.New(session.New(), _getAwsConfig())

	params := &ecs.DescribeTaskDefinitionInput{
		TaskDefinition: aws.String(taskDefinitionName),
	}

	result, err := svc.DescribeTaskDefinition(params)

	if (err != nil) {
		errState(err.Error())
	}

	return result
}

func _stopTask(cluster, taskArn string) {
	svc := ecs.New(session.New(), _getAwsConfig())

	params := &ecs.StopTaskInput{
		Cluster: aws.String(cluster),
		Reason: aws.String("Stopped by WriterTool"),
		Task: aws.String(taskArn),
	}

	_, err := svc.StopTask(params)

	if (err != nil) {
		errState(err.Error())
	}
}

func _waitForNewTask(cluster string, service string, tasks []*string) {
	fmt.Print("  Waiting for new task ")
	newTask := ""

	attempts := 60
	sleepTime := 2
	newTaskAttempts := 5

	for i := 0; i < attempts; i++ {
		currentTasks := _listTasks(cluster, service)
		if (len(currentTasks.TaskArns) < len(tasks) || len(currentTasks.TaskArns) == 0) {
			fmt.Print(".")
		} else if (newTask == "") {
			newTask = _findNewTask(tasks, currentTasks.TaskArns)
			newTaskAttempts--;

			if (newTask == "" && newTaskAttempts >= 0) {
				fmt.Print("?")
			} else if (newTask != "") {
				fmt.Println(" done")
				fmt.Print("  New task: " + newTask + " ")
			} else {
				errState("No new task found among tasks")
			}

		}

		if (newTask != "") {
			taskStates := _describeTasks(cluster, []*string{aws.String(newTask)})
			if (*taskStates.Tasks[0].LastStatus == "RUNNING") {
				fmt.Println(" done")
				return
			} else if (*taskStates.Tasks[0].LastStatus == "STOPPED") {
				fmt.Print("X\n")
				errState("Could not start new task")
			} else {
				fmt.Print(".")
			}
		}

		time.Sleep(time.Duration(sleepTime) * time.Second)
	}

	errStatef("Task [%s] did not start in %d attempts", newTask, attempts)

}

func _findNewTask(tasks, currentTasks []*string) string {
	for i := 0; i < len(tasks); i++ {
		for j := 0; j < len(currentTasks); j++ {
			if *tasks[i] != *currentTasks[j] {
				return *currentTasks[j]
			}
		}
	}

	return ""
}

func GetServiceArn(clusterArn, name string) string {
	clusterArns := _listServices(clusterArn);

	for i := 0; i < len(clusterArns.ServiceArns); i++ {
		arn := clusterArns.ServiceArns[i];
		if (ExtractName(arn) == name) {
			return *arn;
		}
	}

	return "";
}
<|endoftext|>


<commit_before>package regionagogo

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"

	"github.com/Workiva/go-datastructures/augmentedtree"
	"github.com/boltdb/bolt"
	"github.com/golang/geo/s2"
	"github.com/golang/protobuf/proto"
	"github.com/hashicorp/golang-lru"
	"github.com/kpawlik/geojson"
)

const (
	mininumViableLevel = 3 // the minimum cell level we accept
	loopBucket         = "loop"
	coverBucket        = "cover"
)

// GeoSearch provides in memory index and query engine for fences lookup
type GeoSearch struct {
	augmentedtree.Tree
	*bolt.DB
	cache *lru.Cache
	Debug bool
}

type GeoSearchOption func(*geoSearchOptions)

type geoSearchOptions struct {
	maxCachedEntries uint
}

// WithCachedEntries enable an LRU cache default is disabled
func WithCachedEntries(maxCachedEntries uint) GeoSearchOption {
	return func(o *geoSearchOptions) {
		o.maxCachedEntries = maxCachedEntries
	}
}

// NewGeoSearch
func NewGeoSearch(dbpath string, opts ...GeoSearchOption) (*GeoSearch, error) {
	db, err := bolt.Open(dbpath, 0600, nil)
	if err != nil {
		return nil, err
	}
	db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucket([]byte(loopBucket))
		if err != nil {
			return fmt.Errorf("create bucket: %s", err)
		}
		_, err = tx.CreateBucket([]byte(coverBucket))
		if err != nil {
			return fmt.Errorf("create bucket: %s", err)
		}
		return nil
	})
	if err != nil {
		return nil, err
	}

	var geoOpts geoSearchOptions

	for _, opt := range opts {
		opt(&geoOpts)
	}

	gs := &GeoSearch{
		Tree: augmentedtree.New(1),
		DB:   db,
	}

	if geoOpts.maxCachedEntries != 0 {
		cache, err := lru.New(int(geoOpts.maxCachedEntries))
		if err != nil {
			return nil, err
		}
		gs.cache = cache
	}

	return gs, nil
}

// ImportGeoData open the DB and load all cells into the segment tree
func (gs *GeoSearch) ImportGeoData() error {
	var count int
	err := gs.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(coverBucket))
		cur := b.Cursor()

		// load the cell ranges into the tree
		var rc RegionCover
		for k, v := cur.First(); k != nil; k, v = cur.Next() {
			count++
			err := proto.Unmarshal(v, &rc)
			if err != nil {
				return err
			}
			if gs.Debug {
				log.Println("read", rc.Cellunion)
			}

			// read back the loopID from the key
			var id uint64
			buf := bytes.NewBuffer(k)
			err = binary.Read(buf, binary.BigEndian, &id)
			if err != nil {
				return err
			}

			for _, cell := range rc.Cellunion {
				s2interval := &S2Interval{CellID: s2.CellID(cell)}
				intervals := gs.Query(s2interval)
				found := false

				if len(intervals) != 0 {
					for _, existInterval := range intervals {
						if existInterval.LowAtDimension(1) == s2interval.LowAtDimension(1) &&
							existInterval.HighAtDimension(1) == s2interval.HighAtDimension(1) {
							if gs.Debug {
								log.Println("added to existing interval", s2interval, id)
							}
							s2interval.LoopIDs = append(s2interval.LoopIDs, id)
							found = true
							break
						}
					}
				}

				if !found {
					// create new interval with current loop
					s2interval.LoopIDs = []uint64{id}
					gs.Add(s2interval)
					if gs.Debug {
						log.Println("added new interval", s2interval, id)
					}
				}
			}
		}

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

	log.Println("loaded", count, "regions")

	return nil
}

// TODO: refactor as Fence ?
func (gs *GeoSearch) RegionByID(loopID uint64) *Region {
	if gs.cache != nil {
		if val, ok := gs.cache.Get(loopID); ok {
			r, _ := val.(*Region)
			return r
		}
	}
	var rs *RegionStorage
	err := gs.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(loopBucket))
		v := b.Get(itob(loopID))

		var frs RegionStorage
		err := proto.Unmarshal(v, &frs)
		if err != nil {
			return err
		}
		rs = &frs
		return nil
	})
	if err != nil {
		return nil
	}
	r := NewRegionFromStorage(rs)
	if gs.cache != nil && r != nil {
		gs.cache.Add(loopID, r)
	}
	return r
}

// StubbingQuery returns the region for the corresponding lat, lng point
func (gs *GeoSearch) StubbingQuery(lat, lng float64) *Region {
	q := s2.CellIDFromLatLng(s2.LatLngFromDegrees(lat, lng))
	i := &S2Interval{CellID: q}

	if gs.Debug {
		log.Println("lookup", lat, lng, q)
	}
	r := gs.Tree.Query(i)

	var foundRegion *Region

	for _, itv := range r {
		sitv := itv.(*S2Interval)
		if gs.Debug {
			log.Println("found", sitv, sitv.LoopIDs)
		}

		// a region can include a smaller region
		// return only the one that is contained in the other
		for _, loopID := range sitv.LoopIDs {
			region := gs.RegionByID(loopID)
			if region != nil && region.Loop.ContainsPoint(q.Point()) {
				if foundRegion == nil {
					foundRegion = region
					continue
				}
				// we take the 1st vertex of the region.Loop if it is contained in previousLoop
				// region loop  is more precise
				if foundRegion.Loop.ContainsPoint(region.Loop.Vertex(0)) {
					foundRegion = region
				}
			}
		}
	}

	return foundRegion
}

// RectQuery perform rectangular query ur upper right bl bottom left
// return a list of loopID
func (gs *GeoSearch) RectQuery(urlat, urlng, bllat, bllng float64, limit int) (Regions, error) {
	rect := s2.RectFromLatLng(s2.LatLngFromDegrees(bllat, bllng))
	rect = rect.AddPoint(s2.LatLngFromDegrees(urlat, urlng))

	rc := &s2.RegionCoverer{MaxLevel: 30, MaxCells: 1}
	covering := rc.Covering(rect)
	if len(covering) != 1 {
		return nil, errors.New("impossible covering")
	}
	i := &S2Interval{CellID: covering[0]}
	r := gs.Tree.Query(i)

	regions := make(map[uint64]*Region)

	for _, itv := range r {
		sitv := itv.(*S2Interval)
		for _, loopID := range sitv.LoopIDs {
			var region *Region
			if v, ok := regions[loopID]; ok {
				region = v
			} else {
				region = gs.RegionByID(loopID)
			}
			// testing the found loop is actually inside the rect
			// (since we are using only one large cover it may be outside)
			if rect.Contains(region.Loop.RectBound()) {
				regions[loopID] = region
			}
		}
	}

	var res []*Region
	for _,v := range regions {
		res = append(res, v)
	}
	return Regions(res), nil
}

// ImportGeoJSONFile will load a geo json and save the polygons into
// the database for later lookup f
// fields are the properties fields you want to be associated with each fences
func (gs *GeoSearch) ImportGeoJSONFile(r io.Reader, fields []string) error {
	var geo geojson.FeatureCollection

	d := json.NewDecoder(r)
	if err := d.Decode(&geo); err != nil {
		return err
	}

	for _, f := range geo.Features {
		geom, err := f.GetGeometry()
		if err != nil {
			return err
		}

		rc := &s2.RegionCoverer{MinLevel: 1, MaxLevel: 30, MaxCells: 8}

		switch geom.GetType() {
		case "Polygon":
			mp := geom.(*geojson.Polygon)
			for _, p := range mp.Coordinates {
				err := gs.insertPolygon(f, p, rc, fields)
				if err != nil {
					return err
				}
			}
		case "MultiPolygon":
			mp := geom.(*geojson.MultiPolygon)
			// multipolygon
			for _, m := range mp.Coordinates {
				// coordinates polygon
				p := m[0]

				err := gs.insertPolygon(f, p, rc, fields)
				if err != nil {
					return err
				}
			}
		default:
			return errors.New("unknown type")
		}
	}

	return nil
}

func (gs *GeoSearch) insertPolygon(f *geojson.Feature, p geojson.Coordinates, rc *s2.RegionCoverer, fields []string) error {
	// polygon
	// do not add last point in storage (first point is last point)
	points := make([]s2.Point, len(p)-1)

	// For type "MultiPolygon", the "coordinates" member must be an array of Polygon coordinate arrays.
	// "Polygon", the "coordinates" member must be an array of LinearRing coordinate arrays.
	// For Polygons with multiple rings, the first must be the exterior ring and any others must be interior rings or holes.

	for i := 0; i < len(p)-1; i++ {
		ll := s2.LatLngFromDegrees(float64(p[i][1]), float64(p[i][0]))
		points[i] = s2.PointFromLatLng(ll)
	}

	// reverse the slice
	for i := len(points)/2 - 1; i >= 0; i-- {
		opp := len(points) - 1 - i
		points[i], points[opp] = points[opp], points[i]
	}

	// TODO: parallelized region cover
	l := LoopRegionFromPoints(points)

	if l.IsEmpty() || l.IsFull() {
		log.Println("invalid loop", f.Properties)
		return nil
	}

	covering := rc.Covering(l)

	data := make(map[string]string)
	for _, field := range fields {
		if v, ok := f.Properties[field].(string); !ok {
			log.Println("can't find field on", f.Properties)
		} else {
			data[field] = v
		}
	}

	cu := make([]uint64, len(covering))
	var invalidLoop bool

	for i, v := range covering {
		// added a security there if the level is too high it probably means the polygon is bogus
		// this to avoid a large cell to cover everything
		if v.Level() < mininumViableLevel {
			log.Print("cell level too big", v.Level(), data)
			invalidLoop = true
		}
		cu[i] = uint64(v)
	}

	// do not insert big loop
	if invalidLoop {
		return nil
	}

	var cpoints []*CPoint

	for _, p := range points {
		ll := s2.LatLngFromPoint(p)
		cpoints = append(cpoints, &CPoint{Lat: float32(ll.Lat.Degrees()), Lng: float32(ll.Lng.Degrees())})
	}

	rs := &RegionStorage{
		Points: cpoints,
		Data:   data,
	}

	return gs.Update(func(tx *bolt.Tx) error {
		loopB := tx.Bucket([]byte(loopBucket))
		coverBucket := tx.Bucket([]byte(coverBucket))

		id, err := loopB.NextSequence()
		if err != nil {
			return err
		}

		buf, err := proto.Marshal(rs)
		if err != nil {
			return err
		}

		if gs.Debug {
			log.Println("inserted", id, data, cu)
		}

		// convert our loopID to bigendian to be used as key
		k := itob(id)

		err = loopB.Put(k, buf)
		if err != nil {
			return err
		}

		// inserting into cover index using the same key
		bufc, err := proto.Marshal(&RegionCover{Cellunion: cu})
		if err != nil {
			return err
		}

		return coverBucket.Put(k, bufc)
	})
}

// itob returns an 8-byte big endian representation of v.
func itob(v uint64) []byte {
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, v)
	return b
}
<commit_msg>fixing CCW issues<commit_after>package regionagogo

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"

	"github.com/Workiva/go-datastructures/augmentedtree"
	"github.com/boltdb/bolt"
	"github.com/golang/geo/s2"
	"github.com/golang/protobuf/proto"
	"github.com/hashicorp/golang-lru"
	"github.com/kpawlik/geojson"
)

const (
	mininumViableLevel = 3 // the minimum cell level we accept
	loopBucket         = "loop"
	coverBucket        = "cover"
)

// GeoSearch provides in memory index and query engine for fences lookup
type GeoSearch struct {
	augmentedtree.Tree
	*bolt.DB
	cache *lru.Cache
	Debug bool
}

type GeoSearchOption func(*geoSearchOptions)

type geoSearchOptions struct {
	maxCachedEntries uint
}

// WithCachedEntries enable an LRU cache default is disabled
func WithCachedEntries(maxCachedEntries uint) GeoSearchOption {
	return func(o *geoSearchOptions) {
		o.maxCachedEntries = maxCachedEntries
	}
}

// NewGeoSearch
func NewGeoSearch(dbpath string, opts ...GeoSearchOption) (*GeoSearch, error) {
	db, err := bolt.Open(dbpath, 0600, nil)
	if err != nil {
		return nil, err
	}
	db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucket([]byte(loopBucket))
		if err != nil {
			return fmt.Errorf("create bucket: %s", err)
		}
		_, err = tx.CreateBucket([]byte(coverBucket))
		if err != nil {
			return fmt.Errorf("create bucket: %s", err)
		}
		return nil
	})
	if err != nil {
		return nil, err
	}

	var geoOpts geoSearchOptions

	for _, opt := range opts {
		opt(&geoOpts)
	}

	gs := &GeoSearch{
		Tree: augmentedtree.New(1),
		DB:   db,
	}

	if geoOpts.maxCachedEntries != 0 {
		cache, err := lru.New(int(geoOpts.maxCachedEntries))
		if err != nil {
			return nil, err
		}
		gs.cache = cache
	}

	return gs, nil
}

// ImportGeoData open the DB and load all cells into the segment tree
func (gs *GeoSearch) ImportGeoData() error {
	var count int
	err := gs.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(coverBucket))
		cur := b.Cursor()

		// load the cell ranges into the tree
		var rc RegionCover
		for k, v := cur.First(); k != nil; k, v = cur.Next() {
			count++
			err := proto.Unmarshal(v, &rc)
			if err != nil {
				return err
			}
			if gs.Debug {
				log.Println("read", rc.Cellunion)
			}

			// read back the loopID from the key
			var id uint64
			buf := bytes.NewBuffer(k)
			err = binary.Read(buf, binary.BigEndian, &id)
			if err != nil {
				return err
			}

			for _, cell := range rc.Cellunion {
				s2interval := &S2Interval{CellID: s2.CellID(cell)}
				intervals := gs.Query(s2interval)
				found := false

				if len(intervals) != 0 {
					for _, existInterval := range intervals {
						if existInterval.LowAtDimension(1) == s2interval.LowAtDimension(1) &&
							existInterval.HighAtDimension(1) == s2interval.HighAtDimension(1) {
							if gs.Debug {
								log.Println("added to existing interval", s2interval, id)
							}
							s2interval.LoopIDs = append(s2interval.LoopIDs, id)
							found = true
							break
						}
					}
				}

				if !found {
					// create new interval with current loop
					s2interval.LoopIDs = []uint64{id}
					gs.Add(s2interval)
					if gs.Debug {
						log.Println("added new interval", s2interval, id)
					}
				}
			}
		}

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

	log.Println("loaded", count, "regions")

	return nil
}

// TODO: refactor as Fence ?
func (gs *GeoSearch) RegionByID(loopID uint64) *Region {
	if gs.cache != nil {
		if val, ok := gs.cache.Get(loopID); ok {
			r, _ := val.(*Region)
			return r
		}
	}
	var rs *RegionStorage
	err := gs.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(loopBucket))
		v := b.Get(itob(loopID))

		var frs RegionStorage
		err := proto.Unmarshal(v, &frs)
		if err != nil {
			return err
		}
		rs = &frs
		return nil
	})
	if err != nil {
		return nil
	}
	r := NewRegionFromStorage(rs)
	if gs.cache != nil && r != nil {
		gs.cache.Add(loopID, r)
	}
	return r
}

// StubbingQuery returns the region for the corresponding lat, lng point
func (gs *GeoSearch) StubbingQuery(lat, lng float64) *Region {
	q := s2.CellIDFromLatLng(s2.LatLngFromDegrees(lat, lng))
	i := &S2Interval{CellID: q}

	if gs.Debug {
		log.Println("lookup", lat, lng, q)
	}
	r := gs.Tree.Query(i)

	var foundRegion *Region

	for _, itv := range r {
		sitv := itv.(*S2Interval)
		if gs.Debug {
			log.Println("found", sitv, sitv.LoopIDs)
		}

		// a region can include a smaller region
		// return only the one that is contained in the other
		for _, loopID := range sitv.LoopIDs {
			region := gs.RegionByID(loopID)
			if region != nil && region.Loop.ContainsPoint(q.Point()) {
				if foundRegion == nil {
					foundRegion = region
					continue
				}
				// we take the 1st vertex of the region.Loop if it is contained in previousLoop
				// region loop  is more precise
				if foundRegion.Loop.ContainsPoint(region.Loop.Vertex(0)) {
					foundRegion = region
				}
			}
		}
	}

	return foundRegion
}

// RectQuery perform rectangular query ur upper right bl bottom left
// return a list of loopID
func (gs *GeoSearch) RectQuery(urlat, urlng, bllat, bllng float64, limit int) (Regions, error) {
	rect := s2.RectFromLatLng(s2.LatLngFromDegrees(bllat, bllng))
	rect = rect.AddPoint(s2.LatLngFromDegrees(urlat, urlng))

	rc := &s2.RegionCoverer{MaxLevel: 30, MaxCells: 1}
	covering := rc.Covering(rect)
	if len(covering) != 1 {
		return nil, errors.New("impossible covering")
	}
	i := &S2Interval{CellID: covering[0]}
	r := gs.Tree.Query(i)

	regions := make(map[uint64]*Region)

	for _, itv := range r {
		sitv := itv.(*S2Interval)
		for _, loopID := range sitv.LoopIDs {
			var region *Region
			if v, ok := regions[loopID]; ok {
				region = v
			} else {
				region = gs.RegionByID(loopID)
			}
			// testing the found loop is actually inside the rect
			// (since we are using only one large cover it may be outside)
			if rect.Contains(region.Loop.RectBound()) {
				regions[loopID] = region
			}
		}
	}

	var res []*Region
	for _, v := range regions {
		res = append(res, v)
	}
	return Regions(res), nil
}

// ImportGeoJSONFile will load a geo json and save the polygons into
// the database for later lookup f
// fields are the properties fields you want to be associated with each fences
func (gs *GeoSearch) ImportGeoJSONFile(r io.Reader, fields []string) error {
	var geo geojson.FeatureCollection

	d := json.NewDecoder(r)
	if err := d.Decode(&geo); err != nil {
		return err
	}

	for _, f := range geo.Features {
		geom, err := f.GetGeometry()
		if err != nil {
			return err
		}

		rc := &s2.RegionCoverer{MinLevel: 1, MaxLevel: 30, MaxCells: 8}

		switch geom.GetType() {
		case "Polygon":
			mp := geom.(*geojson.Polygon)
			for _, p := range mp.Coordinates {
				err := gs.insertPolygon(f, p, rc, fields)
				if err != nil {
					return err
				}
			}
		case "MultiPolygon":
			mp := geom.(*geojson.MultiPolygon)
			// multipolygon
			for _, m := range mp.Coordinates {
				// coordinates polygon
				p := m[0]

				err := gs.insertPolygon(f, p, rc, fields)
				if err != nil {
					return err
				}
			}
		default:
			return errors.New("unknown type")
		}
	}

	return nil
}

func (gs *GeoSearch) insertPolygon(f *geojson.Feature, p geojson.Coordinates, rc *s2.RegionCoverer, fields []string) error {
	if isClockwisePolygon(p) {
		reversePolygon(p)
	}
	// polygon
	// do not add last point in storage (first point is last point)
	points := make([]s2.Point, len(p)-1)

	// For type "MultiPolygon", the "coordinates" member must be an array of Polygon coordinate arrays.
	// "Polygon", the "coordinates" member must be an array of LinearRing coordinate arrays.
	// For Polygons with multiple rings, the first must be the exterior ring and any others must be interior rings or holes.

	for i := 0; i < len(p)-1; i++ {
		ll := s2.LatLngFromDegrees(float64(p[i][1]), float64(p[i][0]))
		points[i] = s2.PointFromLatLng(ll)
	}

	// TODO: parallelized region cover
	l := LoopRegionFromPoints(points)

	if l.IsEmpty() || l.IsFull() {
		log.Println("invalid loop", f.Properties)
		return nil
	}

	covering := rc.Covering(l)

	data := make(map[string]string)
	for _, field := range fields {
		if v, ok := f.Properties[field].(string); !ok {
			log.Println("can't find field on", f.Properties)
		} else {
			data[field] = v
		}
	}

	cu := make([]uint64, len(covering))
	var invalidLoop bool

	for i, v := range covering {
		// added a security there if the level is too high it probably means the polygon is bogus
		// this to avoid a large cell to cover everything
		if v.Level() < mininumViableLevel {
			log.Print("cell level too big", v.Level(), data)
			invalidLoop = true
		}
		cu[i] = uint64(v)
	}

	// do not insert big loop
	if invalidLoop {
		return nil
	}

	var cpoints []*CPoint

	for _, p := range points {
		ll := s2.LatLngFromPoint(p)
		cpoints = append(cpoints, &CPoint{Lat: float32(ll.Lat.Degrees()), Lng: float32(ll.Lng.Degrees())})
	}

	rs := &RegionStorage{
		Points: cpoints,
		Data:   data,
	}

	return gs.Update(func(tx *bolt.Tx) error {
		loopB := tx.Bucket([]byte(loopBucket))
		coverBucket := tx.Bucket([]byte(coverBucket))

		id, err := loopB.NextSequence()
		if err != nil {
			return err
		}

		buf, err := proto.Marshal(rs)
		if err != nil {
			return err
		}

		if gs.Debug {
			log.Println("inserted", id, data, cu)
		}

		// convert our loopID to bigendian to be used as key
		k := itob(id)

		err = loopB.Put(k, buf)
		if err != nil {
			return err
		}

		// inserting into cover index using the same key
		bufc, err := proto.Marshal(&RegionCover{Cellunion: cu})
		if err != nil {
			return err
		}

		return coverBucket.Put(k, bufc)
	})
}

func isClockwisePolygon(p geojson.Coordinates) bool {
	sum := 0.0
	for i, coord := range p[:len(p)-1] {
		next := p[i+1]
		sum += float64((next[0] - coord[0]) * (next[1] + coord[1]))
	}
	if sum == 0 {
		return true
	}
	return sum > 0
}

func reversePolygon(p geojson.Coordinates) {
	for i := len(p)/2 - 1; i >= 0; i-- {
		opp := len(p) - 1 - i
		p[i], p[opp] = p[opp], p[i]
	}
}

// itob returns an 8-byte big endian representation of v.
func itob(v uint64) []byte {
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, v)
	return b
}
<|endoftext|>


<commit_before>package middleware

// The original work was derived from Goji's middleware, source:
// https://github.com/zenazn/goji/tree/master/web/middleware

import (
	"bufio"
	"io"
	"net"
	"net/http"
)

// WrapResponseWriter is a proxy around an http.ResponseWriter that allows you to hook
// into various parts of the response process.
type WrapResponseWriter interface {
	http.ResponseWriter
	// Status returns the HTTP status of the request, or 0 if one has not
	// yet been sent.
	Status() int
	// BytesWritten returns the total number of bytes sent to the client.
	BytesWritten() int
	// Tee causes the response body to be written to the given io.Writer in
	// addition to proxying the writes through. Only one io.Writer can be
	// tee'd to at once: setting a second one will overwrite the first.
	// Writes will be sent to the proxy before being written to this
	// io.Writer. It is illegal for the tee'd writer to be modified
	// concurrently with writes.
	Tee(io.Writer)
	// Unwrap returns the original proxied target.
	Unwrap() http.ResponseWriter
}

// basicWriter wraps a http.ResponseWriter that implements the minimal
// http.ResponseWriter interface.
type basicWriter struct {
	http.ResponseWriter
	wroteHeader bool
	code        int
	bytes       int
	tee         io.Writer
}

func (b *basicWriter) WriteHeader(code int) {
	if !b.wroteHeader {
		b.code = code
		b.wroteHeader = true
		b.ResponseWriter.WriteHeader(code)
	}
}
func (b *basicWriter) Write(buf []byte) (int, error) {
	b.WriteHeader(http.StatusOK)
	n, err := b.ResponseWriter.Write(buf)
	if b.tee != nil {
		_, err2 := b.tee.Write(buf[:n])
		// Prefer errors generated by the proxied writer.
		if err == nil {
			err = err2
		}
	}
	b.bytes += n
	return n, err
}
func (b *basicWriter) maybeWriteHeader() {
	if !b.wroteHeader {
		b.WriteHeader(http.StatusOK)
	}
}
func (b *basicWriter) Status() int {
	return b.code
}
func (b *basicWriter) BytesWritten() int {
	return b.bytes
}
func (b *basicWriter) Tee(w io.Writer) {
	b.tee = w
}
func (b *basicWriter) Unwrap() http.ResponseWriter {
	return b.ResponseWriter
}

type flushWriter struct {
	basicWriter
}

func (f *flushWriter) Flush() {
	fl := f.basicWriter.ResponseWriter.(http.Flusher)
	fl.Flush()
}

var _ http.Flusher = &flushWriter{}

// httpFancyWriter is a HTTP writer that additionally satisfies http.CloseNotifier,
// http.Flusher, http.Hijacker, and io.ReaderFrom. It exists for the common case
// of wrapping the http.ResponseWriter that package http gives you, in order to
// make the proxied object support the full method set of the proxied object.
type httpFancyWriter struct {
	basicWriter
}

func (f *httpFancyWriter) CloseNotify() <-chan bool {
	cn := f.basicWriter.ResponseWriter.(http.CloseNotifier)
	return cn.CloseNotify()
}
func (f *httpFancyWriter) Flush() {
	fl := f.basicWriter.ResponseWriter.(http.Flusher)
	fl.Flush()
}
func (f *httpFancyWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	hj := f.basicWriter.ResponseWriter.(http.Hijacker)
	return hj.Hijack()
}
func (f *httpFancyWriter) ReadFrom(r io.Reader) (int64, error) {
	if f.basicWriter.tee != nil {
		return io.Copy(&f.basicWriter, r)
	}
	rf := f.basicWriter.ResponseWriter.(io.ReaderFrom)
	f.basicWriter.maybeWriteHeader()
	return rf.ReadFrom(r)
}

var _ http.CloseNotifier = &httpFancyWriter{}
var _ http.Flusher = &httpFancyWriter{}
var _ http.Hijacker = &httpFancyWriter{}
var _ io.ReaderFrom = &httpFancyWriter{}

// http2FancyWriter is a HTTP2 writer that additionally satisfies http.CloseNotifier,
// http.Flusher, and io.ReaderFrom. It exists for the common case
// of wrapping the http.ResponseWriter that package http gives you, in order to
// make the proxied object support the full method set of the proxied object.
type http2FancyWriter struct {
	basicWriter
}

func (f *http2FancyWriter) CloseNotify() <-chan bool {
	cn := f.basicWriter.ResponseWriter.(http.CloseNotifier)
	return cn.CloseNotify()
}
func (f *http2FancyWriter) Flush() {
	fl := f.basicWriter.ResponseWriter.(http.Flusher)
	fl.Flush()
}

var _ http.CloseNotifier = &http2FancyWriter{}
var _ http.Flusher = &http2FancyWriter{}
<commit_msg>Update basicWriter.bytes with ReadFrom result (#234)<commit_after>package middleware

// The original work was derived from Goji's middleware, source:
// https://github.com/zenazn/goji/tree/master/web/middleware

import (
	"bufio"
	"io"
	"net"
	"net/http"
)

// WrapResponseWriter is a proxy around an http.ResponseWriter that allows you to hook
// into various parts of the response process.
type WrapResponseWriter interface {
	http.ResponseWriter
	// Status returns the HTTP status of the request, or 0 if one has not
	// yet been sent.
	Status() int
	// BytesWritten returns the total number of bytes sent to the client.
	BytesWritten() int
	// Tee causes the response body to be written to the given io.Writer in
	// addition to proxying the writes through. Only one io.Writer can be
	// tee'd to at once: setting a second one will overwrite the first.
	// Writes will be sent to the proxy before being written to this
	// io.Writer. It is illegal for the tee'd writer to be modified
	// concurrently with writes.
	Tee(io.Writer)
	// Unwrap returns the original proxied target.
	Unwrap() http.ResponseWriter
}

// basicWriter wraps a http.ResponseWriter that implements the minimal
// http.ResponseWriter interface.
type basicWriter struct {
	http.ResponseWriter
	wroteHeader bool
	code        int
	bytes       int
	tee         io.Writer
}

func (b *basicWriter) WriteHeader(code int) {
	if !b.wroteHeader {
		b.code = code
		b.wroteHeader = true
		b.ResponseWriter.WriteHeader(code)
	}
}
func (b *basicWriter) Write(buf []byte) (int, error) {
	b.WriteHeader(http.StatusOK)
	n, err := b.ResponseWriter.Write(buf)
	if b.tee != nil {
		_, err2 := b.tee.Write(buf[:n])
		// Prefer errors generated by the proxied writer.
		if err == nil {
			err = err2
		}
	}
	b.bytes += n
	return n, err
}
func (b *basicWriter) maybeWriteHeader() {
	if !b.wroteHeader {
		b.WriteHeader(http.StatusOK)
	}
}
func (b *basicWriter) Status() int {
	return b.code
}
func (b *basicWriter) BytesWritten() int {
	return b.bytes
}
func (b *basicWriter) Tee(w io.Writer) {
	b.tee = w
}
func (b *basicWriter) Unwrap() http.ResponseWriter {
	return b.ResponseWriter
}

type flushWriter struct {
	basicWriter
}

func (f *flushWriter) Flush() {
	fl := f.basicWriter.ResponseWriter.(http.Flusher)
	fl.Flush()
}

var _ http.Flusher = &flushWriter{}

// httpFancyWriter is a HTTP writer that additionally satisfies http.CloseNotifier,
// http.Flusher, http.Hijacker, and io.ReaderFrom. It exists for the common case
// of wrapping the http.ResponseWriter that package http gives you, in order to
// make the proxied object support the full method set of the proxied object.
type httpFancyWriter struct {
	basicWriter
}

func (f *httpFancyWriter) CloseNotify() <-chan bool {
	cn := f.basicWriter.ResponseWriter.(http.CloseNotifier)
	return cn.CloseNotify()
}
func (f *httpFancyWriter) Flush() {
	fl := f.basicWriter.ResponseWriter.(http.Flusher)
	fl.Flush()
}
func (f *httpFancyWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	hj := f.basicWriter.ResponseWriter.(http.Hijacker)
	return hj.Hijack()
}
func (f *httpFancyWriter) ReadFrom(r io.Reader) (int64, error) {
	if f.basicWriter.tee != nil {
		return io.Copy(&f.basicWriter, r)
	}
	rf := f.basicWriter.ResponseWriter.(io.ReaderFrom)
	f.basicWriter.maybeWriteHeader()
	n, err := rf.ReadFrom(r)
	f.basicWriter.bytes += int(n)
	return n, err
}

var _ http.CloseNotifier = &httpFancyWriter{}
var _ http.Flusher = &httpFancyWriter{}
var _ http.Hijacker = &httpFancyWriter{}
var _ io.ReaderFrom = &httpFancyWriter{}

// http2FancyWriter is a HTTP2 writer that additionally satisfies http.CloseNotifier,
// http.Flusher, and io.ReaderFrom. It exists for the common case
// of wrapping the http.ResponseWriter that package http gives you, in order to
// make the proxied object support the full method set of the proxied object.
type http2FancyWriter struct {
	basicWriter
}

func (f *http2FancyWriter) CloseNotify() <-chan bool {
	cn := f.basicWriter.ResponseWriter.(http.CloseNotifier)
	return cn.CloseNotify()
}
func (f *http2FancyWriter) Flush() {
	fl := f.basicWriter.ResponseWriter.(http.Flusher)
	fl.Flush()
}

var _ http.CloseNotifier = &http2FancyWriter{}
var _ http.Flusher = &http2FancyWriter{}
<|endoftext|>


<commit_before>package wavefront

import (
	"fmt"
	"net"
	"regexp"
	"sort"
	"strconv"
	"strings"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/outputs"
	"log"
)

type Wavefront struct {
	Prefix          string
	Host            string
	Port            int
	SimpleFields    bool
	MetricSeparator string
	ConvertPaths    bool
	UseRegex        bool
	SourceOverride  []string
	DebugAll        bool
}

// catch many of the invalid chars that could appear in a metric or tag name
var sanitizedChars = strings.NewReplacer(
	"!", "-", "@", "-", "#", "-", "$", "-", "%", "-", "^", "-", "&", "-",
	"*", "-", "(", "-", ")", "-", "+", "-", "`", "-", "'", "-", "\"", "-",
	"[", "-", "]", "-", "{", "-", "}", "-", ":", "-", ";", "-", "<", "-",
	">", "-", ",", "-", "?", "-", "/", "-", "\\", "-", "|", "-", " ", "-",
)

// instead of Replacer which may miss some special characters we can use a regex pattern, but this is significantly slower than Replacer
var sanitizedRegex, _ = regexp.Compile("[^a-zA-Z\\d_.-]")

var tagValueReplacer = strings.NewReplacer("\"", "\\\"", "*", "-")

var pathReplacer = strings.NewReplacer("_", "_")

var sampleConfig = `
  ## prefix for metrics keys
  #prefix = "my.specific.prefix."

  ## DNS name of the wavefront proxy server
  host = "wavefront.example.com"

  ## Port that the Wavefront proxy server listens on
  port = 2878

  ## wether to use "value" for name of simple fields
  #simple_fields = false

  ## character to use between metric and field name.  defaults to . (dot)
  #metric_separator = "."

  ## Convert metric name paths to use metricSeperator character
  ## When true (default) will convert all _ (underscore) chartacters in final metric name
  #convert_paths = true

  ## Use Regex to sanitize metric and tag names from invalid characters
  ## Regex is more thorough, but significantly slower
  #use_regex = false

  ## point tags to use as the source name for Wavefront (if none found, host will be used)
  #source_override = ["hostname", "snmp_host", "node_host"]

  ## Print additional debug information requires debug = true at the agent level
  #debug_all = false
`

type MetricLine struct {
	Metric    string
	Value     string
	Timestamp int64
	Tags      string
}

func (w *Wavefront) Connect() error {

	if w.ConvertPaths && w.MetricSeparator == "_" {
		w.ConvertPaths = false
	}
	if w.ConvertPaths {
		pathReplacer = strings.NewReplacer("_", w.MetricSeparator)
	}

	// Test Connection to Wavefront proxy Server
	uri := fmt.Sprintf("%s:%d", w.Host, w.Port)
	tcpAddr, err := net.ResolveTCPAddr("tcp", uri)
	if err != nil {
		return fmt.Errorf("Wavefront: TCP address cannot be resolved %s", err.Error())
	}
	connection, err := net.DialTCP("tcp", nil, tcpAddr)
	if err != nil {
		return fmt.Errorf("Wavefront: TCP connect fail %s", err.Error())
	}
	defer connection.Close()
	return nil
}

func (w *Wavefront) Write(metrics []telegraf.Metric) error {
	if len(metrics) == 0 {
		return nil
	}

	// Send Data to Wavefront proxy Server
	uri := fmt.Sprintf("%s:%d", w.Host, w.Port)
	tcpAddr, _ := net.ResolveTCPAddr("tcp", uri)
	connection, err := net.DialTCP("tcp", nil, tcpAddr)
	if err != nil {
		return fmt.Errorf("Wavefront: TCP connect fail %s", err.Error())
	}
	defer connection.Close()

	for _, m := range metrics {
		for _, metric := range buildMetrics(m, w) {
			messageLine := fmt.Sprintf("%s %s %v %s\n", metric.Metric, metric.Value, metric.Timestamp, metric.Tags)
			log.Printf("D! Output [wavefront] %s", messageLine)
			_, err := connection.Write([]byte(messageLine))
			if err != nil {
				return fmt.Errorf("Wavefront: TCP writing error %s", err.Error())
			}
		}
	}

	return nil
}

func buildTags(mTags map[string]string, w *Wavefront) []string {
	sourceTagFound := false

	for _, s := range w.SourceOverride {
		for k, v := range mTags {
			if k == s {
				mTags["source"] = v
				mTags["telegraf_host"] = mTags["host"]
				sourceTagFound = true
				delete(mTags, k)
				break
			}
		}
		if sourceTagFound {
			break
		}
	}

	if !sourceTagFound {
		mTags["source"] = mTags["host"]
	}
	delete(mTags, "host")

	tags := make([]string, len(mTags))
	index := 0
	for k, v := range mTags {
		if w.UseRegex {
			tags[index] = fmt.Sprintf("%s=\"%s\"", sanitizedRegex.ReplaceAllString(k, "-"), tagValueReplacer.Replace(v))
		} else {
			tags[index] = fmt.Sprintf("%s=\"%s\"", sanitizedChars.Replace(k), tagValueReplacer.Replace(v))
		}

		index++
	}

	sort.Strings(tags)
	return tags
}

func buildMetrics(m telegraf.Metric, w *Wavefront) []*MetricLine {
	if w.DebugAll {
		log.Printf("D! Output [wavefront] original name: %s\n", m.Name())
	}

	ret := []*MetricLine{}
	for fieldName, value := range m.Fields() {
		if w.DebugAll {
			log.Printf("D! Output [wavefront] original field: %s\n", fieldName)
		}

		var name string
		if !w.SimpleFields && fieldName == "value" {
			name = fmt.Sprintf("%s%s", w.Prefix, m.Name())
		} else {
			name = fmt.Sprintf("%s%s%s%s", w.Prefix, m.Name(), w.MetricSeparator, fieldName)
		}

		if w.UseRegex {
			name = sanitizedRegex.ReplaceAllLiteralString(name, "-")
		} else {
			name = sanitizedChars.Replace(name)
		}

		if w.ConvertPaths {
			name = pathReplacer.Replace(name)
		}

		metric := &MetricLine{
			Metric:    name,
			Timestamp: m.UnixNano() / 1000000000,
		}
		metricValue, buildError := buildValue(value, metric.Metric)
		if buildError != nil {
			log.Printf("D! Output [wavefront] %s\n", buildError.Error())
			continue
		}
		metric.Value = metricValue
		tagsSlice := buildTags(m.Tags(), w)
		metric.Tags = fmt.Sprint(strings.Join(tagsSlice, " "))
		ret = append(ret, metric)
	}
	return ret
}

func buildValue(v interface{}, name string) (string, error) {
	var retv string
	switch p := v.(type) {
	case int64:
		retv = IntToString(int64(p))
	case uint64:
		retv = UIntToString(uint64(p))
	case float64:
		retv = FloatToString(float64(p))
	default:
		return retv, fmt.Errorf("unexpected type: %T, with value: %v, for: %s", v, v, name)
	}
	return retv, nil
}

func IntToString(input_num int64) string {
	return strconv.FormatInt(input_num, 10)
}

func UIntToString(input_num uint64) string {
	return strconv.FormatUint(input_num, 10)
}

func FloatToString(input_num float64) string {
	return strconv.FormatFloat(input_num, 'f', 6, 64)
}

func (w *Wavefront) SampleConfig() string {
	return sampleConfig
}

func (w *Wavefront) Description() string {
	return "Configuration for Wavefront server to send metrics to"
}

func (w *Wavefront) Close() error {
	return nil
}

func init() {
	outputs.Add("wavefront", func() telegraf.Output {
		return &Wavefront{
			MetricSeparator: ".",
			ConvertPaths:    true,
		}
	})
}
<commit_msg>Do not remove the wavefront tag when it becomes the source<commit_after>package wavefront

import (
	"fmt"
	"net"
	"regexp"
	"sort"
	"strconv"
	"strings"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/outputs"
	"log"
)

type Wavefront struct {
	Prefix          string
	Host            string
	Port            int
	SimpleFields    bool
	MetricSeparator string
	ConvertPaths    bool
	UseRegex        bool
	SourceOverride  []string
	DebugAll        bool
}

// catch many of the invalid chars that could appear in a metric or tag name
var sanitizedChars = strings.NewReplacer(
	"!", "-", "@", "-", "#", "-", "$", "-", "%", "-", "^", "-", "&", "-",
	"*", "-", "(", "-", ")", "-", "+", "-", "`", "-", "'", "-", "\"", "-",
	"[", "-", "]", "-", "{", "-", "}", "-", ":", "-", ";", "-", "<", "-",
	">", "-", ",", "-", "?", "-", "/", "-", "\\", "-", "|", "-", " ", "-",
)

// instead of Replacer which may miss some special characters we can use a regex pattern, but this is significantly slower than Replacer
var sanitizedRegex, _ = regexp.Compile("[^a-zA-Z\\d_.-]")

var tagValueReplacer = strings.NewReplacer("\"", "\\\"", "*", "-")

var pathReplacer = strings.NewReplacer("_", "_")

var sampleConfig = `
  ## prefix for metrics keys
  #prefix = "my.specific.prefix."

  ## DNS name of the wavefront proxy server
  host = "wavefront.example.com"

  ## Port that the Wavefront proxy server listens on
  port = 2878

  ## wether to use "value" for name of simple fields
  #simple_fields = false

  ## character to use between metric and field name.  defaults to . (dot)
  #metric_separator = "."

  ## Convert metric name paths to use metricSeperator character
  ## When true (default) will convert all _ (underscore) chartacters in final metric name
  #convert_paths = true

  ## Use Regex to sanitize metric and tag names from invalid characters
  ## Regex is more thorough, but significantly slower
  #use_regex = false

  ## point tags to use as the source name for Wavefront (if none found, host will be used)
  #source_override = ["hostname", "snmp_host", "node_host"]

  ## Print additional debug information requires debug = true at the agent level
  #debug_all = false
`

type MetricLine struct {
	Metric    string
	Value     string
	Timestamp int64
	Tags      string
}

func (w *Wavefront) Connect() error {

	if w.ConvertPaths && w.MetricSeparator == "_" {
		w.ConvertPaths = false
	}
	if w.ConvertPaths {
		pathReplacer = strings.NewReplacer("_", w.MetricSeparator)
	}

	// Test Connection to Wavefront proxy Server
	uri := fmt.Sprintf("%s:%d", w.Host, w.Port)
	tcpAddr, err := net.ResolveTCPAddr("tcp", uri)
	if err != nil {
		return fmt.Errorf("Wavefront: TCP address cannot be resolved %s", err.Error())
	}
	connection, err := net.DialTCP("tcp", nil, tcpAddr)
	if err != nil {
		return fmt.Errorf("Wavefront: TCP connect fail %s", err.Error())
	}
	defer connection.Close()
	return nil
}

func (w *Wavefront) Write(metrics []telegraf.Metric) error {
	if len(metrics) == 0 {
		return nil
	}

	// Send Data to Wavefront proxy Server
	uri := fmt.Sprintf("%s:%d", w.Host, w.Port)
	tcpAddr, _ := net.ResolveTCPAddr("tcp", uri)
	connection, err := net.DialTCP("tcp", nil, tcpAddr)
	if err != nil {
		return fmt.Errorf("Wavefront: TCP connect fail %s", err.Error())
	}
	defer connection.Close()

	for _, m := range metrics {
		for _, metric := range buildMetrics(m, w) {
			messageLine := fmt.Sprintf("%s %s %v %s\n", metric.Metric, metric.Value, metric.Timestamp, metric.Tags)
			log.Printf("D! Output [wavefront] %s", messageLine)
			_, err := connection.Write([]byte(messageLine))
			if err != nil {
				return fmt.Errorf("Wavefront: TCP writing error %s", err.Error())
			}
		}
	}

	return nil
}

func buildTags(mTags map[string]string, w *Wavefront) []string {
	sourceTagFound := false

	for _, s := range w.SourceOverride {
		for k, v := range mTags {
			if k == s {
				mTags["source"] = v
				mTags["telegraf_host"] = mTags["host"]
				sourceTagFound = true
				// Commenting this out because we want the source tag to also appear as a tag
				// delete(mTags, k)
				break
			}
		}
		if sourceTagFound {
			break
		}
	}

	if !sourceTagFound {
		mTags["source"] = mTags["host"]
	}
	delete(mTags, "host")

	tags := make([]string, len(mTags))
	index := 0
	for k, v := range mTags {
		if w.UseRegex {
			tags[index] = fmt.Sprintf("%s=\"%s\"", sanitizedRegex.ReplaceAllString(k, "-"), tagValueReplacer.Replace(v))
		} else {
			tags[index] = fmt.Sprintf("%s=\"%s\"", sanitizedChars.Replace(k), tagValueReplacer.Replace(v))
		}

		index++
	}

	sort.Strings(tags)
	return tags
}

func buildMetrics(m telegraf.Metric, w *Wavefront) []*MetricLine {
	if w.DebugAll {
		log.Printf("D! Output [wavefront] original name: %s\n", m.Name())
	}

	ret := []*MetricLine{}
	for fieldName, value := range m.Fields() {
		if w.DebugAll {
			log.Printf("D! Output [wavefront] original field: %s\n", fieldName)
		}

		var name string
		if !w.SimpleFields && fieldName == "value" {
			name = fmt.Sprintf("%s%s", w.Prefix, m.Name())
		} else {
			name = fmt.Sprintf("%s%s%s%s", w.Prefix, m.Name(), w.MetricSeparator, fieldName)
		}

		if w.UseRegex {
			name = sanitizedRegex.ReplaceAllLiteralString(name, "-")
		} else {
			name = sanitizedChars.Replace(name)
		}

		if w.ConvertPaths {
			name = pathReplacer.Replace(name)
		}

		metric := &MetricLine{
			Metric:    name,
			Timestamp: m.UnixNano() / 1000000000,
		}
		metricValue, buildError := buildValue(value, metric.Metric)
		if buildError != nil {
			log.Printf("D! Output [wavefront] %s\n", buildError.Error())
			continue
		}
		metric.Value = metricValue
		tagsSlice := buildTags(m.Tags(), w)
		metric.Tags = fmt.Sprint(strings.Join(tagsSlice, " "))
		ret = append(ret, metric)
	}
	return ret
}

func buildValue(v interface{}, name string) (string, error) {
	var retv string
	switch p := v.(type) {
	case int64:
		retv = IntToString(int64(p))
	case uint64:
		retv = UIntToString(uint64(p))
	case float64:
		retv = FloatToString(float64(p))
	default:
		return retv, fmt.Errorf("unexpected type: %T, with value: %v, for: %s", v, v, name)
	}
	return retv, nil
}

func IntToString(input_num int64) string {
	return strconv.FormatInt(input_num, 10)
}

func UIntToString(input_num uint64) string {
	return strconv.FormatUint(input_num, 10)
}

func FloatToString(input_num float64) string {
	return strconv.FormatFloat(input_num, 'f', 6, 64)
}

func (w *Wavefront) SampleConfig() string {
	return sampleConfig
}

func (w *Wavefront) Description() string {
	return "Configuration for Wavefront server to send metrics to"
}

func (w *Wavefront) Close() error {
	return nil
}

func init() {
	outputs.Add("wavefront", func() telegraf.Output {
		return &Wavefront{
			MetricSeparator: ".",
			ConvertPaths:    true,
		}
	})
}
<|endoftext|>