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

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

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

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

package api_test

import (
	"encoding/json"
	"math/rand"
	"reflect"
	"testing"

	"github.com/davecgh/go-spew/spew"
	flag "github.com/spf13/pflag"
	"github.com/ugorji/go/codec"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/meta"
	"k8s.io/kubernetes/pkg/api/testapi"
	apitesting "k8s.io/kubernetes/pkg/api/testing"
	"k8s.io/kubernetes/pkg/api/unversioned"
	"k8s.io/kubernetes/pkg/api/v1"
	_ "k8s.io/kubernetes/pkg/apis/extensions"
	_ "k8s.io/kubernetes/pkg/apis/extensions/v1beta1"
	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/util"
	"k8s.io/kubernetes/pkg/util/sets"
)

var fuzzIters = flag.Int("fuzz-iters", 20, "How many fuzzing iterations to do.")

var codecsToTest = []func(version unversioned.GroupVersion, item runtime.Object) (runtime.Codec, error){
	func(version unversioned.GroupVersion, item runtime.Object) (runtime.Codec, error) {
		return testapi.GetCodecForObject(item)
	},
}

func fuzzInternalObject(t *testing.T, forVersion unversioned.GroupVersion, item runtime.Object, seed int64) runtime.Object {
	apitesting.FuzzerFor(t, forVersion, rand.NewSource(seed)).Fuzz(item)

	j, err := meta.TypeAccessor(item)
	if err != nil {
		t.Fatalf("Unexpected error %v for %#v", err, item)
	}
	j.SetKind("")
	j.SetAPIVersion("")

	return item
}

func roundTrip(t *testing.T, codec runtime.Codec, item runtime.Object) {
	printer := spew.ConfigState{DisableMethods: true}

	gvk, err := api.Scheme.ObjectKind(item)
	t.Logf("fully qualified kind for %v is %v with codec %v", reflect.TypeOf(item), gvk, codec)

	name := reflect.TypeOf(item).Elem().Name()
	data, err := runtime.Encode(codec, item)
	if err != nil {
		t.Errorf("%v: %v (%s)", name, err, printer.Sprintf("%#v", item))
		return
	}

	obj2, err := runtime.Decode(codec, data)
	if err != nil {
		t.Errorf("0: %v: %v\nCodec: %v\nData: %s\nSource: %#v", name, err, codec, string(data), printer.Sprintf("%#v", item))
		return
	}
	if !api.Semantic.DeepEqual(item, obj2) {
		t.Errorf("\n1: %v: diff: %v\nCodec: %v\nSource:\n\n%#v\n\nEncoded:\n\n%s\n\nFinal:\n\n%#v", name, util.ObjectGoPrintDiff(item, obj2), codec, printer.Sprintf("%#v", item), string(data), printer.Sprintf("%#v", obj2))
		return
	}

	obj3 := reflect.New(reflect.TypeOf(item).Elem()).Interface().(runtime.Object)
	err = runtime.DecodeInto(codec, data, obj3)
	if err != nil {
		t.Errorf("2: %v: %v", name, err)
		return
	}
	if !api.Semantic.DeepEqual(item, obj3) {
		t.Errorf("3: %v: diff: %v\nCodec: %v", name, util.ObjectDiff(item, obj3), codec)
		return
	}
}

// roundTripSame verifies the same source object is tested in all API versions.
func roundTripSame(t *testing.T, item runtime.Object, except ...string) {
	set := sets.NewString(except...)
	seed := rand.Int63()
	fuzzInternalObject(t, testapi.Default.InternalGroupVersion(), item, seed)

	version := *testapi.Default.GroupVersion()
	codecs := []runtime.Codec{}
	for _, fn := range codecsToTest {
		codec, err := fn(version, item)
		if err != nil {
			t.Errorf("unable to get codec: %v", err)
			return
		}
		codecs = append(codecs, codec)
	}

	if !set.Has(version.String()) {
		fuzzInternalObject(t, version, item, seed)
		for _, codec := range codecs {
			roundTrip(t, codec, item)
		}
	}
}

// For debugging problems
func TestSpecificKind(t *testing.T) {
	api.Scheme.Log(t)
	defer api.Scheme.Log(nil)

	kind := "Pod"
	for i := 0; i < *fuzzIters; i++ {
		doRoundTripTest(kind, t)
		if t.Failed() {
			break
		}
	}
}

func TestList(t *testing.T) {
	api.Scheme.Log(t)
	defer api.Scheme.Log(nil)

	kind := "List"
	item, err := api.Scheme.New(api.SchemeGroupVersion.WithKind(kind))
	if err != nil {
		t.Errorf("Couldn't make a %v? %v", kind, err)
		return
	}
	roundTripSame(t, item)
}

var nonRoundTrippableTypes = sets.NewString("ExportOptions")

var nonInternalRoundTrippableTypes = sets.NewString("List", "ListOptions", "ExportOptions")
var nonRoundTrippableTypesByVersion = map[string][]string{}

func TestRoundTripTypes(t *testing.T) {
	// api.Scheme.Log(t)
	// defer api.Scheme.Log(nil)

	for kind := range api.Scheme.KnownTypes(testapi.Default.InternalGroupVersion()) {
		t.Logf("working on %v in %v", kind, testapi.Default.InternalGroupVersion())
		if nonRoundTrippableTypes.Has(kind) {
			continue
		}
		// Try a few times, since runTest uses random values.
		for i := 0; i < *fuzzIters; i++ {
			doRoundTripTest(kind, t)
			if t.Failed() {
				break
			}
		}
	}
}

func doRoundTripTest(kind string, t *testing.T) {
	item, err := api.Scheme.New(testapi.Default.InternalGroupVersion().WithKind(kind))
	if err != nil {
		t.Fatalf("Couldn't make a %v? %v", kind, err)
	}
	if _, err := meta.TypeAccessor(item); err != nil {
		t.Fatalf("%q is not a TypeMeta and cannot be tested - add it to nonRoundTrippableTypes: %v", kind, err)
	}
	if api.Scheme.Recognizes(testapi.Default.GroupVersion().WithKind(kind)) {
		roundTripSame(t, item, nonRoundTrippableTypesByVersion[kind]...)
	}
	if !nonInternalRoundTrippableTypes.Has(kind) {
		roundTrip(t, api.Codec, fuzzInternalObject(t, testapi.Default.InternalGroupVersion(), item, rand.Int63()))
	}
}

func TestEncode_Ptr(t *testing.T) {
	grace := int64(30)
	pod := &api.Pod{
		ObjectMeta: api.ObjectMeta{
			Labels: map[string]string{"name": "foo"},
		},
		Spec: api.PodSpec{
			RestartPolicy: api.RestartPolicyAlways,
			DNSPolicy:     api.DNSClusterFirst,

			TerminationGracePeriodSeconds: &grace,

			SecurityContext: &api.PodSecurityContext{},
		},
	}
	obj := runtime.Object(pod)
	data, err := testapi.Default.Codec().Encode(obj)
	obj2, err2 := testapi.Default.Codec().Decode(data)
	if err != nil || err2 != nil {
		t.Fatalf("Failure: '%v' '%v'", err, err2)
	}
	if _, ok := obj2.(*api.Pod); !ok {
		t.Fatalf("Got wrong type")
	}
	if !api.Semantic.DeepEqual(obj2, pod) {
		t.Errorf("\nExpected:\n\n %#v,\n\nGot:\n\n %#vDiff: %v\n\n", pod, obj2, util.ObjectDiff(obj2, pod))

	}
}

func TestBadJSONRejection(t *testing.T) {
	badJSONMissingKind := []byte(`{ }`)
	if _, err := testapi.Default.Codec().Decode(badJSONMissingKind); err == nil {
		t.Errorf("Did not reject despite lack of kind field: %s", badJSONMissingKind)
	}
	badJSONUnknownType := []byte(`{"kind": "bar"}`)
	if _, err1 := testapi.Default.Codec().Decode(badJSONUnknownType); err1 == nil {
		t.Errorf("Did not reject despite use of unknown type: %s", badJSONUnknownType)
	}
	/*badJSONKindMismatch := []byte(`{"kind": "Pod"}`)
	if err2 := DecodeInto(badJSONKindMismatch, &Minion{}); err2 == nil {
		t.Errorf("Kind is set but doesn't match the object type: %s", badJSONKindMismatch)
	}*/
}

func TestUnversionedTypes(t *testing.T) {
	testcases := []runtime.Object{
		&unversioned.Status{Status: "Failure", Message: "something went wrong"},
		&unversioned.APIVersions{Versions: []string{"A", "B", "C"}},
		&unversioned.APIGroupList{Groups: []unversioned.APIGroup{{Name: "mygroup"}}},
		&unversioned.APIGroup{Name: "mygroup"},
		&unversioned.APIResourceList{GroupVersion: "mygroup/myversion"},
	}

	for _, obj := range testcases {
		// Make sure the unversioned codec can encode
		unversionedJSON, err := api.Codec.Encode(obj)
		if err != nil {
			t.Errorf("%v: unexpected error: %v", obj, err)
			continue
		}

		// Make sure the versioned codec under test can decode
		versionDecodedObject, err := testapi.Default.Codec().Decode(unversionedJSON)
		if err != nil {
			t.Errorf("%v: unexpected error: %v", obj, err)
			continue
		}
		// Make sure it decodes correctly
		if !reflect.DeepEqual(obj, versionDecodedObject) {
			t.Errorf("%v: expected %#v, got %#v", obj, obj, versionDecodedObject)
			continue
		}
	}
}

const benchmarkSeed = 100

func benchmarkItems() []v1.Pod {
	apiObjectFuzzer := apitesting.FuzzerFor(nil, api.SchemeGroupVersion, rand.NewSource(benchmarkSeed))
	items := make([]v1.Pod, 2)
	for i := range items {
		apiObjectFuzzer.Fuzz(&items[i])
	}
	return items
}

// BenchmarkEncodeCodec measures the cost of performing a codec encode, which includes
// reflection (to clear APIVersion and Kind)
func BenchmarkEncodeCodec(b *testing.B) {
	items := benchmarkItems()
	width := len(items)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if _, err := testapi.Default.Codec().Encode(&items[i%width]); err != nil {
			b.Fatal(err)
		}
	}
	b.StopTimer()
}

// BenchmarkEncodeJSONMarshal provides a baseline for regular JSON encode performance
func BenchmarkEncodeJSONMarshal(b *testing.B) {
	items := benchmarkItems()
	width := len(items)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if _, err := json.Marshal(&items[i%width]); err != nil {
			b.Fatal(err)
		}
	}
	b.StopTimer()
}

func BenchmarkDecodeCodec(b *testing.B) {
	codec := testapi.Default.Codec()
	items := benchmarkItems()
	width := len(items)
	encoded := make([][]byte, width)
	for i := range items {
		data, err := runtime.Encode(codec, &items[i])
		if err != nil {
			b.Fatal(err)
		}
		encoded[i] = data
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if _, err := runtime.Decode(codec, encoded[i%width]); err != nil {
			b.Fatal(err)
		}
	}
	b.StopTimer()
}

func BenchmarkDecodeIntoCodec(b *testing.B) {
	codec := testapi.Default.Codec()
	items := benchmarkItems()
	width := len(items)
	encoded := make([][]byte, width)
	for i := range items {
		data, err := runtime.Encode(codec, &items[i])
		if err != nil {
			b.Fatal(err)
		}
		encoded[i] = data
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		obj := v1.Pod{}
		if err := runtime.DecodeInto(codec, encoded[i%width], &obj); err != nil {
			b.Fatal(err)
		}
	}
	b.StopTimer()
}

// BenchmarkDecodeJSON provides a baseline for regular JSON decode performance
func BenchmarkDecodeIntoJSON(b *testing.B) {
	codec := testapi.Default.Codec()
	items := benchmarkItems()
	width := len(items)
	encoded := make([][]byte, width)
	for i := range items {
		data, err := runtime.Encode(codec, &items[i])
		if err != nil {
			b.Fatal(err)
		}
		encoded[i] = data
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		obj := v1.Pod{}
		if err := json.Unmarshal(encoded[i%width], &obj); err != nil {
			b.Fatal(err)
		}
	}
	b.StopTimer()
}

// BenchmarkDecodeJSON provides a baseline for codecgen JSON decode performance
func BenchmarkDecodeIntoJSONCodecGen(b *testing.B) {
	kcodec := testapi.Default.Codec()
	items := benchmarkItems()
	width := len(items)
	encoded := make([][]byte, width)
	for i := range items {
		data, err := runtime.Encode(kcodec, &items[i])
		if err != nil {
			b.Fatal(err)
		}
		encoded[i] = data
	}
	handler := &codec.JsonHandle{}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		obj := v1.Pod{}
		if err := codec.NewDecoderBytes(encoded[i%width], handler).Decode(&obj); err != nil {
			b.Fatal(err)
		}
	}
	b.StopTimer()
}
<commit_msg>Get rid of rubbish logs from serialization_test<commit_after>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

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

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

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

package api_test

import (
	"encoding/json"
	"math/rand"
	"reflect"
	"testing"

	"github.com/davecgh/go-spew/spew"
	flag "github.com/spf13/pflag"
	"github.com/ugorji/go/codec"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/meta"
	"k8s.io/kubernetes/pkg/api/testapi"
	apitesting "k8s.io/kubernetes/pkg/api/testing"
	"k8s.io/kubernetes/pkg/api/unversioned"
	"k8s.io/kubernetes/pkg/api/v1"
	_ "k8s.io/kubernetes/pkg/apis/extensions"
	_ "k8s.io/kubernetes/pkg/apis/extensions/v1beta1"
	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/util"
	"k8s.io/kubernetes/pkg/util/sets"
)

var fuzzIters = flag.Int("fuzz-iters", 20, "How many fuzzing iterations to do.")

var codecsToTest = []func(version unversioned.GroupVersion, item runtime.Object) (runtime.Codec, error){
	func(version unversioned.GroupVersion, item runtime.Object) (runtime.Codec, error) {
		return testapi.GetCodecForObject(item)
	},
}

func fuzzInternalObject(t *testing.T, forVersion unversioned.GroupVersion, item runtime.Object, seed int64) runtime.Object {
	apitesting.FuzzerFor(t, forVersion, rand.NewSource(seed)).Fuzz(item)

	j, err := meta.TypeAccessor(item)
	if err != nil {
		t.Fatalf("Unexpected error %v for %#v", err, item)
	}
	j.SetKind("")
	j.SetAPIVersion("")

	return item
}

func roundTrip(t *testing.T, codec runtime.Codec, item runtime.Object) {
	printer := spew.ConfigState{DisableMethods: true}

	name := reflect.TypeOf(item).Elem().Name()
	data, err := runtime.Encode(codec, item)
	if err != nil {
		t.Errorf("%v: %v (%s)", name, err, printer.Sprintf("%#v", item))
		return
	}

	obj2, err := runtime.Decode(codec, data)
	if err != nil {
		t.Errorf("0: %v: %v\nCodec: %v\nData: %s\nSource: %#v", name, err, codec, string(data), printer.Sprintf("%#v", item))
		return
	}
	if !api.Semantic.DeepEqual(item, obj2) {
		t.Errorf("\n1: %v: diff: %v\nCodec: %v\nSource:\n\n%#v\n\nEncoded:\n\n%s\n\nFinal:\n\n%#v", name, util.ObjectGoPrintDiff(item, obj2), codec, printer.Sprintf("%#v", item), string(data), printer.Sprintf("%#v", obj2))
		return
	}

	obj3 := reflect.New(reflect.TypeOf(item).Elem()).Interface().(runtime.Object)
	err = runtime.DecodeInto(codec, data, obj3)
	if err != nil {
		t.Errorf("2: %v: %v", name, err)
		return
	}
	if !api.Semantic.DeepEqual(item, obj3) {
		t.Errorf("3: %v: diff: %v\nCodec: %v", name, util.ObjectDiff(item, obj3), codec)
		return
	}
}

// roundTripSame verifies the same source object is tested in all API versions.
func roundTripSame(t *testing.T, item runtime.Object, except ...string) {
	set := sets.NewString(except...)
	seed := rand.Int63()
	fuzzInternalObject(t, testapi.Default.InternalGroupVersion(), item, seed)

	version := *testapi.Default.GroupVersion()
	codecs := []runtime.Codec{}
	for _, fn := range codecsToTest {
		codec, err := fn(version, item)
		if err != nil {
			t.Errorf("unable to get codec: %v", err)
			return
		}
		codecs = append(codecs, codec)
	}

	if !set.Has(version.String()) {
		fuzzInternalObject(t, version, item, seed)
		for _, codec := range codecs {
			roundTrip(t, codec, item)
		}
	}
}

// For debugging problems
func TestSpecificKind(t *testing.T) {
	// api.Scheme.Log(t)
	// defer api.Scheme.Log(nil)

	kind := "Pod"
	for i := 0; i < *fuzzIters; i++ {
		doRoundTripTest(kind, t)
		if t.Failed() {
			break
		}
	}
}

func TestList(t *testing.T) {
	// api.Scheme.Log(t)
	// defer api.Scheme.Log(nil)

	kind := "List"
	item, err := api.Scheme.New(api.SchemeGroupVersion.WithKind(kind))
	if err != nil {
		t.Errorf("Couldn't make a %v? %v", kind, err)
		return
	}
	roundTripSame(t, item)
}

var nonRoundTrippableTypes = sets.NewString("ExportOptions")

var nonInternalRoundTrippableTypes = sets.NewString("List", "ListOptions", "ExportOptions")
var nonRoundTrippableTypesByVersion = map[string][]string{}

func TestRoundTripTypes(t *testing.T) {
	// api.Scheme.Log(t)
	// defer api.Scheme.Log(nil)

	for kind := range api.Scheme.KnownTypes(testapi.Default.InternalGroupVersion()) {
		t.Logf("working on %v in %v", kind, testapi.Default.InternalGroupVersion())
		if nonRoundTrippableTypes.Has(kind) {
			continue
		}
		// Try a few times, since runTest uses random values.
		for i := 0; i < *fuzzIters; i++ {
			doRoundTripTest(kind, t)
			if t.Failed() {
				break
			}
		}
	}
}

func doRoundTripTest(kind string, t *testing.T) {
	item, err := api.Scheme.New(testapi.Default.InternalGroupVersion().WithKind(kind))
	if err != nil {
		t.Fatalf("Couldn't make a %v? %v", kind, err)
	}
	if _, err := meta.TypeAccessor(item); err != nil {
		t.Fatalf("%q is not a TypeMeta and cannot be tested - add it to nonRoundTrippableTypes: %v", kind, err)
	}
	if api.Scheme.Recognizes(testapi.Default.GroupVersion().WithKind(kind)) {
		roundTripSame(t, item, nonRoundTrippableTypesByVersion[kind]...)
	}
	if !nonInternalRoundTrippableTypes.Has(kind) {
		roundTrip(t, api.Codec, fuzzInternalObject(t, testapi.Default.InternalGroupVersion(), item, rand.Int63()))
	}
}

func TestEncode_Ptr(t *testing.T) {
	grace := int64(30)
	pod := &api.Pod{
		ObjectMeta: api.ObjectMeta{
			Labels: map[string]string{"name": "foo"},
		},
		Spec: api.PodSpec{
			RestartPolicy: api.RestartPolicyAlways,
			DNSPolicy:     api.DNSClusterFirst,

			TerminationGracePeriodSeconds: &grace,

			SecurityContext: &api.PodSecurityContext{},
		},
	}
	obj := runtime.Object(pod)
	data, err := testapi.Default.Codec().Encode(obj)
	obj2, err2 := testapi.Default.Codec().Decode(data)
	if err != nil || err2 != nil {
		t.Fatalf("Failure: '%v' '%v'", err, err2)
	}
	if _, ok := obj2.(*api.Pod); !ok {
		t.Fatalf("Got wrong type")
	}
	if !api.Semantic.DeepEqual(obj2, pod) {
		t.Errorf("\nExpected:\n\n %#v,\n\nGot:\n\n %#vDiff: %v\n\n", pod, obj2, util.ObjectDiff(obj2, pod))

	}
}

func TestBadJSONRejection(t *testing.T) {
	badJSONMissingKind := []byte(`{ }`)
	if _, err := testapi.Default.Codec().Decode(badJSONMissingKind); err == nil {
		t.Errorf("Did not reject despite lack of kind field: %s", badJSONMissingKind)
	}
	badJSONUnknownType := []byte(`{"kind": "bar"}`)
	if _, err1 := testapi.Default.Codec().Decode(badJSONUnknownType); err1 == nil {
		t.Errorf("Did not reject despite use of unknown type: %s", badJSONUnknownType)
	}
	/*badJSONKindMismatch := []byte(`{"kind": "Pod"}`)
	if err2 := DecodeInto(badJSONKindMismatch, &Minion{}); err2 == nil {
		t.Errorf("Kind is set but doesn't match the object type: %s", badJSONKindMismatch)
	}*/
}

func TestUnversionedTypes(t *testing.T) {
	testcases := []runtime.Object{
		&unversioned.Status{Status: "Failure", Message: "something went wrong"},
		&unversioned.APIVersions{Versions: []string{"A", "B", "C"}},
		&unversioned.APIGroupList{Groups: []unversioned.APIGroup{{Name: "mygroup"}}},
		&unversioned.APIGroup{Name: "mygroup"},
		&unversioned.APIResourceList{GroupVersion: "mygroup/myversion"},
	}

	for _, obj := range testcases {
		// Make sure the unversioned codec can encode
		unversionedJSON, err := api.Codec.Encode(obj)
		if err != nil {
			t.Errorf("%v: unexpected error: %v", obj, err)
			continue
		}

		// Make sure the versioned codec under test can decode
		versionDecodedObject, err := testapi.Default.Codec().Decode(unversionedJSON)
		if err != nil {
			t.Errorf("%v: unexpected error: %v", obj, err)
			continue
		}
		// Make sure it decodes correctly
		if !reflect.DeepEqual(obj, versionDecodedObject) {
			t.Errorf("%v: expected %#v, got %#v", obj, obj, versionDecodedObject)
			continue
		}
	}
}

const benchmarkSeed = 100

func benchmarkItems() []v1.Pod {
	apiObjectFuzzer := apitesting.FuzzerFor(nil, api.SchemeGroupVersion, rand.NewSource(benchmarkSeed))
	items := make([]v1.Pod, 2)
	for i := range items {
		apiObjectFuzzer.Fuzz(&items[i])
	}
	return items
}

// BenchmarkEncodeCodec measures the cost of performing a codec encode, which includes
// reflection (to clear APIVersion and Kind)
func BenchmarkEncodeCodec(b *testing.B) {
	items := benchmarkItems()
	width := len(items)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if _, err := testapi.Default.Codec().Encode(&items[i%width]); err != nil {
			b.Fatal(err)
		}
	}
	b.StopTimer()
}

// BenchmarkEncodeJSONMarshal provides a baseline for regular JSON encode performance
func BenchmarkEncodeJSONMarshal(b *testing.B) {
	items := benchmarkItems()
	width := len(items)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if _, err := json.Marshal(&items[i%width]); err != nil {
			b.Fatal(err)
		}
	}
	b.StopTimer()
}

func BenchmarkDecodeCodec(b *testing.B) {
	codec := testapi.Default.Codec()
	items := benchmarkItems()
	width := len(items)
	encoded := make([][]byte, width)
	for i := range items {
		data, err := runtime.Encode(codec, &items[i])
		if err != nil {
			b.Fatal(err)
		}
		encoded[i] = data
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if _, err := runtime.Decode(codec, encoded[i%width]); err != nil {
			b.Fatal(err)
		}
	}
	b.StopTimer()
}

func BenchmarkDecodeIntoCodec(b *testing.B) {
	codec := testapi.Default.Codec()
	items := benchmarkItems()
	width := len(items)
	encoded := make([][]byte, width)
	for i := range items {
		data, err := runtime.Encode(codec, &items[i])
		if err != nil {
			b.Fatal(err)
		}
		encoded[i] = data
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		obj := v1.Pod{}
		if err := runtime.DecodeInto(codec, encoded[i%width], &obj); err != nil {
			b.Fatal(err)
		}
	}
	b.StopTimer()
}

// BenchmarkDecodeJSON provides a baseline for regular JSON decode performance
func BenchmarkDecodeIntoJSON(b *testing.B) {
	codec := testapi.Default.Codec()
	items := benchmarkItems()
	width := len(items)
	encoded := make([][]byte, width)
	for i := range items {
		data, err := runtime.Encode(codec, &items[i])
		if err != nil {
			b.Fatal(err)
		}
		encoded[i] = data
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		obj := v1.Pod{}
		if err := json.Unmarshal(encoded[i%width], &obj); err != nil {
			b.Fatal(err)
		}
	}
	b.StopTimer()
}

// BenchmarkDecodeJSON provides a baseline for codecgen JSON decode performance
func BenchmarkDecodeIntoJSONCodecGen(b *testing.B) {
	kcodec := testapi.Default.Codec()
	items := benchmarkItems()
	width := len(items)
	encoded := make([][]byte, width)
	for i := range items {
		data, err := runtime.Encode(kcodec, &items[i])
		if err != nil {
			b.Fatal(err)
		}
		encoded[i] = data
	}
	handler := &codec.JsonHandle{}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		obj := v1.Pod{}
		if err := codec.NewDecoderBytes(encoded[i%width], handler).Decode(&obj); err != nil {
			b.Fatal(err)
		}
	}
	b.StopTimer()
}
<|endoftext|>


<commit_before>// Copyright 2020 TiKV Project 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package autoscaling

import (
	"context"
	"fmt"
	"math"
	"net"
	"strings"
	"time"

	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	promClient "github.com/prometheus/client_golang/api"
	promAPI "github.com/prometheus/client_golang/api/prometheus/v1"
	promModel "github.com/prometheus/common/model"
	"github.com/tikv/pd/pkg/errs"
	"go.uber.org/zap"
)

const (
	tikvSumCPUUsageMetricsPattern = `sum(increase(tikv_thread_cpu_seconds_total[%s])) by (instance, kubernetes_namespace)`
	tidbSumCPUUsageMetricsPattern = `sum(increase(process_cpu_seconds_total{job="tidb"}[%s])) by (instance, kubernetes_namespace)`
	tikvCPUQuotaMetricsPattern    = `tikv_server_cpu_cores_quota`
	tidbCPUQuotaMetricsPattern    = `tidb_server_maxprocs`
	instanceLabelName             = "instance"
	namespaceLabelName            = "kubernetes_namespace"
	addressFormat                 = "pod-name.peer-svc.namespace.svc:port"

	httpRequestTimeout = 5 * time.Second
)

// PrometheusQuerier query metrics from Prometheus
type PrometheusQuerier struct {
	api promAPI.API
}

// NewPrometheusQuerier returns a PrometheusQuerier
func NewPrometheusQuerier(client promClient.Client) *PrometheusQuerier {
	return &PrometheusQuerier{
		api: promAPI.NewAPI(client),
	}
}

type promQLBuilderFn func(*QueryOptions) (string, error)

var queryBuilderFnMap = map[MetricType]promQLBuilderFn{
	CPUQuota: buildCPUQuotaPromQL,
	CPUUsage: buildCPUUsagePromQL,
}

// Query do the real query on Prometheus and returns metric value for each instance
func (prom *PrometheusQuerier) Query(options *QueryOptions) (QueryResult, error) {
	builderFn, ok := queryBuilderFnMap[options.metric]
	if !ok {
		return nil, errs.ErrUnsupportedMetricsType.FastGenByArgs(options.metric)
	}

	query, err := builderFn(options)
	if err != nil {
		return nil, err
	}

	resp, err := prom.queryMetricsFromPrometheus(query, options.timestamp)
	if err != nil {
		return nil, err
	}

	result, err := extractInstancesFromResponse(resp, options.addresses)
	if err != nil {
		return nil, err
	}

	return result, nil
}

func (prom *PrometheusQuerier) queryMetricsFromPrometheus(query string, timestamp time.Time) (promModel.Value, error) {
	ctx, cancel := context.WithTimeout(context.Background(), httpRequestTimeout)
	defer cancel()

	resp, warnings, err := prom.api.Query(ctx, query, timestamp)

	if err != nil {
		return nil, errs.ErrPrometheusQuery.Wrap(err).FastGenWithCause()
	}

	if warnings != nil && len(warnings) > 0 {
		log.Warn("prometheus query returns with warnings", zap.Strings("warnings", warnings))
	}

	return resp, nil
}

func extractInstancesFromResponse(resp promModel.Value, addresses []string) (QueryResult, error) {
	if resp == nil {
		return nil, errs.ErrEmptyMetricsResponse.FastGenByArgs()
	}

	if resp.Type() != promModel.ValVector {
		return nil, errs.ErrUnexpectedType.FastGenByArgs(resp.Type().String())
	}

	vector, ok := resp.(promModel.Vector)

	if !ok {
		return nil, errs.ErrTypeConversion.FastGenByArgs()
	}

	if len(vector) == 0 {
		return nil, errs.ErrEmptyMetricsResult.FastGenByArgs("query metrics duration must be at least twice the Prometheus scrape interval")
	}

	instancesSet := map[string]string{}
	for _, addr := range addresses {
		instanceName, err := getInstanceNameFromAddress(addr)
		if err == nil {
			instancesSet[instanceName] = addr
		}
	}

	result := make(QueryResult)

	for _, sample := range vector {
		podName, ok := sample.Metric[instanceLabelName]
		if !ok {
			continue
		}

		namespace, ok := sample.Metric[namespaceLabelName]
		if !ok {
			continue
		}

		instanceName := buildInstanceIdentifier(string(podName), string(namespace))

		if addr, ok := instancesSet[instanceName]; ok {
			result[addr] = float64(sample.Value)
		}
	}

	return result, nil
}

var cpuUsagePromQLTemplate = map[ComponentType]string{
	TiDB: tidbSumCPUUsageMetricsPattern,
	TiKV: tikvSumCPUUsageMetricsPattern,
}

var cpuQuotaPromQLTemplate = map[ComponentType]string{
	TiDB: tidbCPUQuotaMetricsPattern,
	TiKV: tikvCPUQuotaMetricsPattern,
}

func buildCPUQuotaPromQL(options *QueryOptions) (string, error) {
	pattern, ok := cpuQuotaPromQLTemplate[options.component]
	if !ok {
		return "", errs.ErrUnsupportedComponentType.FastGenByArgs(options.component)
	}

	query := pattern
	return query, nil
}

func buildCPUUsagePromQL(options *QueryOptions) (string, error) {
	pattern, ok := cpuUsagePromQLTemplate[options.component]
	if !ok {
		return "", errs.ErrUnsupportedComponentType.FastGenByArgs(options.component)
	}

	query := fmt.Sprintf(pattern, getDurationExpression(options.duration))
	return query, nil
}

// this function assumes that addr is already a valid resolvable address
// returns in format "podname_namespace"
func getInstanceNameFromAddress(addr string) (string, error) {
	// In K8s, a StatefulSet pod address is composed of pod-name.peer-svc.namespace.svc:port
	// Extract the hostname part without port
	hostname := addr
	portColonIdx := strings.LastIndex(addr, ":")
	if portColonIdx >= 0 {
		hostname = addr[:portColonIdx]
	}

	// Just to make sure it is not an IP address
	ip := net.ParseIP(hostname)
	if ip != nil {
		// Hostname is an IP address, return the whole address
		return "", errors.Errorf("address %s is an ip address", addr)
	}

	parts := strings.Split(hostname, ".")
	if len(parts) < 4 {
		return "", errors.Errorf("address %s does not match the expected format %s", addr, addressFormat)
	}

	podName, namespace := parts[0], parts[2]

	return buildInstanceIdentifier(podName, namespace), nil
}

func buildInstanceIdentifier(podName string, namespace string) string {
	return fmt.Sprintf("%s_%s", podName, namespace)
}

func getDurationExpression(duration time.Duration) string {
	// Prometheus only accept single unit duration like 10s, 2m
	// and the time.Duration.String() method returns the duration like 2m0s, 2m30s,
	// so we need to express the duration in seconds like 120s
	seconds := int64(math.Floor(duration.Seconds()))
	return fmt.Sprintf("%ds", seconds)
}
<commit_msg>tidbSumCPUUsageMetricsPattern problem (#3421) (#3421)<commit_after>// Copyright 2020 TiKV Project 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package autoscaling

import (
	"context"
	"fmt"
	"math"
	"net"
	"strings"
	"time"

	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	promClient "github.com/prometheus/client_golang/api"
	promAPI "github.com/prometheus/client_golang/api/prometheus/v1"
	promModel "github.com/prometheus/common/model"
	"github.com/tikv/pd/pkg/errs"
	"go.uber.org/zap"
)

const (
	tikvSumCPUUsageMetricsPattern = `sum(increase(tikv_thread_cpu_seconds_total[%s])) by (instance, kubernetes_namespace)`
	tidbSumCPUUsageMetricsPattern = `sum(increase(process_cpu_seconds_total{component="tidb"}[%s])) by (instance, kubernetes_namespace)`
	tikvCPUQuotaMetricsPattern    = `tikv_server_cpu_cores_quota`
	tidbCPUQuotaMetricsPattern    = `tidb_server_maxprocs`
	instanceLabelName             = "instance"
	namespaceLabelName            = "kubernetes_namespace"
	addressFormat                 = "pod-name.peer-svc.namespace.svc:port"

	httpRequestTimeout = 5 * time.Second
)

// PrometheusQuerier query metrics from Prometheus
type PrometheusQuerier struct {
	api promAPI.API
}

// NewPrometheusQuerier returns a PrometheusQuerier
func NewPrometheusQuerier(client promClient.Client) *PrometheusQuerier {
	return &PrometheusQuerier{
		api: promAPI.NewAPI(client),
	}
}

type promQLBuilderFn func(*QueryOptions) (string, error)

var queryBuilderFnMap = map[MetricType]promQLBuilderFn{
	CPUQuota: buildCPUQuotaPromQL,
	CPUUsage: buildCPUUsagePromQL,
}

// Query do the real query on Prometheus and returns metric value for each instance
func (prom *PrometheusQuerier) Query(options *QueryOptions) (QueryResult, error) {
	builderFn, ok := queryBuilderFnMap[options.metric]
	if !ok {
		return nil, errs.ErrUnsupportedMetricsType.FastGenByArgs(options.metric)
	}

	query, err := builderFn(options)
	if err != nil {
		return nil, err
	}

	resp, err := prom.queryMetricsFromPrometheus(query, options.timestamp)
	if err != nil {
		return nil, err
	}

	result, err := extractInstancesFromResponse(resp, options.addresses)
	if err != nil {
		return nil, err
	}

	return result, nil
}

func (prom *PrometheusQuerier) queryMetricsFromPrometheus(query string, timestamp time.Time) (promModel.Value, error) {
	ctx, cancel := context.WithTimeout(context.Background(), httpRequestTimeout)
	defer cancel()

	resp, warnings, err := prom.api.Query(ctx, query, timestamp)

	if err != nil {
		return nil, errs.ErrPrometheusQuery.Wrap(err).FastGenWithCause()
	}

	if warnings != nil && len(warnings) > 0 {
		log.Warn("prometheus query returns with warnings", zap.Strings("warnings", warnings))
	}

	return resp, nil
}

func extractInstancesFromResponse(resp promModel.Value, addresses []string) (QueryResult, error) {
	if resp == nil {
		return nil, errs.ErrEmptyMetricsResponse.FastGenByArgs()
	}

	if resp.Type() != promModel.ValVector {
		return nil, errs.ErrUnexpectedType.FastGenByArgs(resp.Type().String())
	}

	vector, ok := resp.(promModel.Vector)

	if !ok {
		return nil, errs.ErrTypeConversion.FastGenByArgs()
	}

	if len(vector) == 0 {
		return nil, errs.ErrEmptyMetricsResult.FastGenByArgs("query metrics duration must be at least twice the Prometheus scrape interval")
	}

	instancesSet := map[string]string{}
	for _, addr := range addresses {
		instanceName, err := getInstanceNameFromAddress(addr)
		if err == nil {
			instancesSet[instanceName] = addr
		}
	}

	result := make(QueryResult)

	for _, sample := range vector {
		podName, ok := sample.Metric[instanceLabelName]
		if !ok {
			continue
		}

		namespace, ok := sample.Metric[namespaceLabelName]
		if !ok {
			continue
		}

		instanceName := buildInstanceIdentifier(string(podName), string(namespace))

		if addr, ok := instancesSet[instanceName]; ok {
			result[addr] = float64(sample.Value)
		}
	}

	return result, nil
}

var cpuUsagePromQLTemplate = map[ComponentType]string{
	TiDB: tidbSumCPUUsageMetricsPattern,
	TiKV: tikvSumCPUUsageMetricsPattern,
}

var cpuQuotaPromQLTemplate = map[ComponentType]string{
	TiDB: tidbCPUQuotaMetricsPattern,
	TiKV: tikvCPUQuotaMetricsPattern,
}

func buildCPUQuotaPromQL(options *QueryOptions) (string, error) {
	pattern, ok := cpuQuotaPromQLTemplate[options.component]
	if !ok {
		return "", errs.ErrUnsupportedComponentType.FastGenByArgs(options.component)
	}

	query := pattern
	return query, nil
}

func buildCPUUsagePromQL(options *QueryOptions) (string, error) {
	pattern, ok := cpuUsagePromQLTemplate[options.component]
	if !ok {
		return "", errs.ErrUnsupportedComponentType.FastGenByArgs(options.component)
	}

	query := fmt.Sprintf(pattern, getDurationExpression(options.duration))
	return query, nil
}

// this function assumes that addr is already a valid resolvable address
// returns in format "podname_namespace"
func getInstanceNameFromAddress(addr string) (string, error) {
	// In K8s, a StatefulSet pod address is composed of pod-name.peer-svc.namespace.svc:port
	// Extract the hostname part without port
	hostname := addr
	portColonIdx := strings.LastIndex(addr, ":")
	if portColonIdx >= 0 {
		hostname = addr[:portColonIdx]
	}

	// Just to make sure it is not an IP address
	ip := net.ParseIP(hostname)
	if ip != nil {
		// Hostname is an IP address, return the whole address
		return "", errors.Errorf("address %s is an ip address", addr)
	}

	parts := strings.Split(hostname, ".")
	if len(parts) < 4 {
		return "", errors.Errorf("address %s does not match the expected format %s", addr, addressFormat)
	}

	podName, namespace := parts[0], parts[2]

	return buildInstanceIdentifier(podName, namespace), nil
}

func buildInstanceIdentifier(podName string, namespace string) string {
	return fmt.Sprintf("%s_%s", podName, namespace)
}

func getDurationExpression(duration time.Duration) string {
	// Prometheus only accept single unit duration like 10s, 2m
	// and the time.Duration.String() method returns the duration like 2m0s, 2m30s,
	// so we need to express the duration in seconds like 120s
	seconds := int64(math.Floor(duration.Seconds()))
	return fmt.Sprintf("%ds", seconds)
}
<|endoftext|>


<commit_before>package stage

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/moby/buildkit/frontend/dockerfile/instructions"
	"github.com/moby/buildkit/frontend/dockerfile/parser"
	"github.com/moby/buildkit/frontend/dockerfile/shell"

	"github.com/flant/werf/pkg/git_repo"
	"github.com/flant/werf/pkg/git_repo/ls_tree"
	"github.com/flant/werf/pkg/git_repo/status"
	"github.com/flant/werf/pkg/image"
	"github.com/flant/werf/pkg/path_matcher"
	"github.com/flant/werf/pkg/util"

	"github.com/flant/logboek"
)

func GenerateDockerfileStage(dockerRunArgs *DockerRunArgs, dockerStages *DockerStages, contextChecksum *ContextChecksum, baseStageOptions *NewBaseStageOptions) *DockerfileStage {
	return newDockerfileStage(dockerRunArgs, dockerStages, contextChecksum, baseStageOptions)
}

func newDockerfileStage(dockerRunArgs *DockerRunArgs, dockerStages *DockerStages, contextChecksum *ContextChecksum, baseStageOptions *NewBaseStageOptions) *DockerfileStage {
	s := &DockerfileStage{}
	s.DockerRunArgs = dockerRunArgs
	s.DockerStages = dockerStages
	s.ContextChecksum = contextChecksum
	s.BaseStage = newBaseStage(Dockerfile, baseStageOptions)

	return s
}

type DockerfileStage struct {
	*DockerRunArgs
	*DockerStages
	*ContextChecksum
	*BaseStage
}

func NewDockerRunArgs(dockerfilePath, target, context string, buildArgs map[string]interface{}, addHost []string) *DockerRunArgs {
	return &DockerRunArgs{
		dockerfilePath: dockerfilePath,
		target:         target,
		context:        context,
		buildArgs:      buildArgs,
		addHost:        addHost,
	}
}

type DockerRunArgs struct {
	dockerfilePath string
	target         string
	context        string
	buildArgs      map[string]interface{}
	addHost        []string
}

func NewDockerStages(dockerStages []instructions.Stage, dockerArgsHash map[string]string, dockerTargetStageIndex int) *DockerStages {
	return &DockerStages{
		dockerStages:           dockerStages,
		dockerTargetStageIndex: dockerTargetStageIndex,
		dockerArgsHash:         dockerArgsHash,
	}
}

type DockerStages struct {
	dockerStages           []instructions.Stage
	dockerArgsHash         map[string]string
	dockerTargetStageIndex int
}

func NewContextChecksum(projectPath string, dockerignorePathMatcher *path_matcher.DockerfileIgnorePathMatcher, localGitRepo *git_repo.Local) *ContextChecksum {
	return &ContextChecksum{
		projectPath:             projectPath,
		dockerignorePathMatcher: dockerignorePathMatcher,
		localGitRepo:            localGitRepo,
	}
}

type ContextChecksum struct {
	projectPath             string
	localGitRepo            *git_repo.Local
	dockerignorePathMatcher *path_matcher.DockerfileIgnorePathMatcher

	mainLsTreeResult *ls_tree.Result
	mainStatusResult *status.Result
}

type dockerfileInstructionInterface interface {
	String() string
	Name() string
}

func (s *DockerfileStage) GetDependencies(_ Conveyor, _, _ image.ImageInterface) (string, error) {
	var dockerMetaArgsString []string
	for key, value := range s.dockerArgsHash {
		dockerMetaArgsString = append(dockerMetaArgsString, fmt.Sprintf("%s=%s", key, value))
	}

	shlex := shell.NewLex(parser.DefaultEscapeToken)

	var stagesDependencies [][]string
	for _, stage := range s.dockerStages {
		var dependencies []string

		dependencies = append(dependencies, s.addHost...)

		resolvedBaseName, err := shlex.ProcessWord(stage.BaseName, dockerMetaArgsString)
		if err != nil {
			return "", err
		}

		dependencies = append(dependencies, resolvedBaseName)

		for _, cmd := range stage.Commands {
			switch c := cmd.(type) {
			case *instructions.ArgCommand:
				dependencies = append(dependencies, c.String())
				if argValue, exist := s.dockerArgsHash[c.Key]; exist {
					dependencies = append(dependencies, argValue)
				}
			case *instructions.AddCommand:
				dependencies = append(dependencies, c.String())

				checksum, err := s.calculateFilesChecksum(c.SourcesAndDest.Sources())
				if err != nil {
					return "", err
				}
				dependencies = append(dependencies, checksum)
			case *instructions.CopyCommand:
				dependencies = append(dependencies, c.String())
				if c.From == "" {
					checksum, err := s.calculateFilesChecksum(c.SourcesAndDest.Sources())
					if err != nil {
						return "", err
					}
					dependencies = append(dependencies, checksum)
				}
			case dockerfileInstructionInterface:
				dependencies = append(dependencies, c.String())
			default:
				panic("runtime error")
			}
		}

		stagesDependencies = append(stagesDependencies, dependencies)
	}

	for ind, stage := range s.dockerStages {
		for relatedStageIndex, relatedStage := range s.dockerStages {
			if ind == relatedStageIndex {
				continue
			}

			if stage.BaseName == relatedStage.Name {
				stagesDependencies[ind] = append(stagesDependencies[ind], stagesDependencies[relatedStageIndex]...)
			}
		}

		for _, cmd := range stage.Commands {
			switch c := cmd.(type) {
			case *instructions.CopyCommand:
				if c.From != "" {
					relatedStageIndex, err := strconv.Atoi(c.From)
					if err != nil || relatedStageIndex >= len(stagesDependencies) {
						logboek.LogWarnF("WARNING: COPY --from with nonexistent dockerfile stage %s detected\n", c.From)
						continue
					} else {
						stagesDependencies[ind] = append(stagesDependencies[ind], stagesDependencies[relatedStageIndex]...)
					}
				}
			}
		}
	}

	return util.Sha256Hash(stagesDependencies[s.dockerTargetStageIndex]...), nil
}

func (s *DockerfileStage) PrepareImage(c Conveyor, prevBuiltImage, img image.ImageInterface) error {
	img.DockerfileImageBuilder().AppendBuildArgs(s.DockerBuildArgs()...)
	return nil
}

func (s *DockerfileStage) DockerBuildArgs() []string {
	var result []string

	if s.dockerfilePath != "" {
		result = append(result, fmt.Sprintf("--file=%s", s.dockerfilePath))
	}

	if s.target != "" {
		result = append(result, fmt.Sprintf("--target=%s", s.target))
	}

	if len(s.buildArgs) != 0 {
		for key, value := range s.buildArgs {
			result = append(result, fmt.Sprintf("--build-arg=%s=%v", key, value))
		}
	}

	for _, addHost := range s.addHost {
		result = append(result, fmt.Sprintf("--add-host=%s", addHost))
	}

	result = append(result, s.context)

	return result
}

func (s *DockerfileStage) calculateFilesChecksum(wildcards []string) (string, error) {
	var checksum string
	var err error

	logProcessMsg := fmt.Sprintf("Calculating files checksum (%v)", wildcards)
	logboek.Debug.LogProcessStart(logProcessMsg, logboek.LevelLogProcessStartOptions{})
	if s.localGitRepo != nil {
		checksum, err = s.calculateFilesChecksumWithLsTree(wildcards)
	} else {
		checksum, err = s.calculateFilesChecksumWithFilesRead(wildcards)
	}

	if err != nil {
		logboek.Debug.LogProcessFail(logboek.LevelLogProcessFailOptions{})
		return "", err
	}

	logboek.Debug.LogProcessEnd(logboek.LevelLogProcessEndOptions{})

	logboek.Debug.LogF("Result checksum: %s\n", checksum)
	logboek.Debug.LogOptionalLn()

	return checksum, nil
}

func (s *DockerfileStage) calculateFilesChecksumWithLsTree(wildcards []string) (string, error) {
	if s.mainLsTreeResult == nil {
		processMsg := fmt.Sprintf("ls-tree (%s)", s.dockerignorePathMatcher.String())
		logboek.Debug.LogProcessStart(processMsg, logboek.LevelLogProcessStartOptions{})
		result, err := s.localGitRepo.LsTree(s.dockerignorePathMatcher)
		if err != nil {
			if err.Error() == "entry not found" {
				logboek.Debug.LogFWithCustomStyle(
					logboek.StyleByName(logboek.FailStyleName),
					"Entry %s is not found\n",
					s.dockerignorePathMatcher.BaseFilepath(),
				)
				logboek.Debug.LogProcessEnd(logboek.LevelLogProcessEndOptions{})
				goto entryNotFoundInGitRepository
			}

			logboek.Debug.LogProcessFail(logboek.LevelLogProcessFailOptions{})
			return "", err
		}
		logboek.Debug.LogProcessEnd(logboek.LevelLogProcessEndOptions{})

		s.mainLsTreeResult = result
	}

entryNotFoundInGitRepository:
	wildcardsPathMatcher := path_matcher.NewSimplePathMatcher(s.dockerignorePathMatcher.BaseFilepath(), wildcards, false)

	var lsTreeResultChecksum string
	if s.mainLsTreeResult != nil {
		blockMsg := fmt.Sprintf("ls-tree (%s)", wildcardsPathMatcher.String())
		logboek.Debug.LogProcessStart(blockMsg, logboek.LevelLogProcessStartOptions{})
		lsTreeResult, err := s.mainLsTreeResult.LsTree(wildcardsPathMatcher)
		if err != nil {
			logboek.Debug.LogProcessFail(logboek.LevelLogProcessFailOptions{})
			return "", err
		}
		logboek.Debug.LogProcessEnd(logboek.LevelLogProcessEndOptions{})

		if !lsTreeResult.IsEmpty() {
			blockMsg = fmt.Sprintf("ls-tree result checksum (%s)", wildcardsPathMatcher.String())
			_ = logboek.Debug.LogBlock(blockMsg, logboek.LevelLogBlockOptions{}, func() error {
				lsTreeResultChecksum = lsTreeResult.Checksum()
				logboek.Debug.LogLn()
				logboek.Debug.LogLn(lsTreeResultChecksum)

				return nil
			})
		}
	}

	if s.mainStatusResult == nil {
		processMsg := fmt.Sprintf("status (%s)", s.dockerignorePathMatcher.String())
		logboek.Debug.LogProcessStart(processMsg, logboek.LevelLogProcessStartOptions{})
		result, err := s.localGitRepo.Status(s.dockerignorePathMatcher)
		if err != nil {
			logboek.Debug.LogProcessFail(logboek.LevelLogProcessFailOptions{})
			return "", err
		}
		logboek.Debug.LogProcessEnd(logboek.LevelLogProcessEndOptions{})

		s.mainStatusResult = result
	}

	blockMsg := fmt.Sprintf("status (%s)", wildcardsPathMatcher.String())
	logboek.Debug.LogProcessStart(blockMsg, logboek.LevelLogProcessStartOptions{})
	statusResult, err := s.mainStatusResult.Status(wildcardsPathMatcher)
	if err != nil {
		logboek.Debug.LogProcessFail(logboek.LevelLogProcessFailOptions{})
		return "", err
	}
	logboek.Debug.LogProcessEnd(logboek.LevelLogProcessEndOptions{})

	var statusResultChecksum string
	if !statusResult.IsEmpty() {
		blockMsg = fmt.Sprintf("Status result checksum (%s)", wildcardsPathMatcher.String())
		_ = logboek.Debug.LogBlock(blockMsg, logboek.LevelLogBlockOptions{}, func() error {
			statusResultChecksum = statusResult.Checksum()
			logboek.Debug.LogLn()
			logboek.Debug.LogLn(statusResultChecksum)
			return nil
		})
	}

	resultChecksum := util.Sha256Hash(lsTreeResultChecksum, statusResultChecksum)

	return resultChecksum, nil
}

func (s *DockerfileStage) calculateFilesChecksumWithFilesRead(wildcards []string) (string, error) {
	var dependencies []string

	for _, wildcard := range wildcards {
		contextWildcard := filepath.Join(s.context, wildcard)

		matches, err := filepath.Glob(contextWildcard)
		if err != nil {
			return "", fmt.Errorf("glob %s failed: %s", contextWildcard, err)
		}

		var fileList []string
		for _, match := range matches {
			matchFileList, err := getAllFiles(match)
			if err != nil {
				return "", fmt.Errorf("walk %s failed: %s", match, err)
			}

			fileList = append(fileList, matchFileList...)
		}

		var finalFileList []string
		for _, filePath := range fileList {
			relFilePath, err := filepath.Rel(s.projectPath, filePath)
			if err != nil {
				panic(fmt.Sprintf("unexpected condition: %s", err))
			} else if relFilePath == "." || relFilePath == ".." || strings.HasPrefix(relFilePath, ".."+string(os.PathSeparator)) {
				panic(fmt.Sprintf("unexpected condition: %s", relFilePath))
			}

			if s.dockerignorePathMatcher.MatchPath(relFilePath) {
				finalFileList = append(finalFileList, filePath)
			}
		}

		for _, filePath := range finalFileList {
			data, err := ioutil.ReadFile(filePath)
			if err != nil {
				return "", fmt.Errorf("read file %s failed: %s", filePath, err)
			}

			dependencies = append(dependencies, string(data))
			logboek.Debug.LogF("File was added: %s\n", strings.TrimPrefix(filePath, s.projectPath+string(os.PathSeparator)))
		}
	}

	checksum := util.Sha256Hash(dependencies...)

	return checksum, nil
}

func getAllFiles(target string) ([]string, error) {
	var fileList []string
	err := filepath.Walk(target, func(path string, f os.FileInfo, err error) error {
		if f.IsDir() {
			return nil
		}

		if f.Mode()&os.ModeSymlink != 0 {
			linkTo, err := os.Readlink(path)
			if err != nil {
				return err
			}

			linkFilePath := filepath.Join(filepath.Dir(path), linkTo)
			exist, err := util.FileExists(linkFilePath)
			if err != nil {
				return err
			} else if !exist {
				return nil
			} else {
				lfinfo, err := os.Stat(linkFilePath)
				if err != nil {
					return err
				}

				if lfinfo.IsDir() {
					// infinite loop detector
					if target == linkFilePath {
						return nil
					}

					lfileList, err := getAllFiles(linkFilePath)
					if err != nil {
						return err
					}

					fileList = append(fileList, lfileList...)
				} else {
					fileList = append(fileList, linkFilePath)
				}

				return nil
			}
		}

		fileList = append(fileList, path)
		return err
	})

	return fileList, err
}
<commit_msg>[dockerfile] Fail build if dockerfile COPY instruction refers to nonexistent stage<commit_after>package stage

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/moby/buildkit/frontend/dockerfile/instructions"
	"github.com/moby/buildkit/frontend/dockerfile/parser"
	"github.com/moby/buildkit/frontend/dockerfile/shell"

	"github.com/flant/werf/pkg/git_repo"
	"github.com/flant/werf/pkg/git_repo/ls_tree"
	"github.com/flant/werf/pkg/git_repo/status"
	"github.com/flant/werf/pkg/image"
	"github.com/flant/werf/pkg/path_matcher"
	"github.com/flant/werf/pkg/util"

	"github.com/flant/logboek"
)

func GenerateDockerfileStage(dockerRunArgs *DockerRunArgs, dockerStages *DockerStages, contextChecksum *ContextChecksum, baseStageOptions *NewBaseStageOptions) *DockerfileStage {
	return newDockerfileStage(dockerRunArgs, dockerStages, contextChecksum, baseStageOptions)
}

func newDockerfileStage(dockerRunArgs *DockerRunArgs, dockerStages *DockerStages, contextChecksum *ContextChecksum, baseStageOptions *NewBaseStageOptions) *DockerfileStage {
	s := &DockerfileStage{}
	s.DockerRunArgs = dockerRunArgs
	s.DockerStages = dockerStages
	s.ContextChecksum = contextChecksum
	s.BaseStage = newBaseStage(Dockerfile, baseStageOptions)

	return s
}

type DockerfileStage struct {
	*DockerRunArgs
	*DockerStages
	*ContextChecksum
	*BaseStage
}

func NewDockerRunArgs(dockerfilePath, target, context string, buildArgs map[string]interface{}, addHost []string) *DockerRunArgs {
	return &DockerRunArgs{
		dockerfilePath: dockerfilePath,
		target:         target,
		context:        context,
		buildArgs:      buildArgs,
		addHost:        addHost,
	}
}

type DockerRunArgs struct {
	dockerfilePath string
	target         string
	context        string
	buildArgs      map[string]interface{}
	addHost        []string
}

func NewDockerStages(dockerStages []instructions.Stage, dockerArgsHash map[string]string, dockerTargetStageIndex int) *DockerStages {
	return &DockerStages{
		dockerStages:           dockerStages,
		dockerTargetStageIndex: dockerTargetStageIndex,
		dockerArgsHash:         dockerArgsHash,
	}
}

type DockerStages struct {
	dockerStages           []instructions.Stage
	dockerArgsHash         map[string]string
	dockerTargetStageIndex int
}

func NewContextChecksum(projectPath string, dockerignorePathMatcher *path_matcher.DockerfileIgnorePathMatcher, localGitRepo *git_repo.Local) *ContextChecksum {
	return &ContextChecksum{
		projectPath:             projectPath,
		dockerignorePathMatcher: dockerignorePathMatcher,
		localGitRepo:            localGitRepo,
	}
}

type ContextChecksum struct {
	projectPath             string
	localGitRepo            *git_repo.Local
	dockerignorePathMatcher *path_matcher.DockerfileIgnorePathMatcher

	mainLsTreeResult *ls_tree.Result
	mainStatusResult *status.Result
}

type dockerfileInstructionInterface interface {
	String() string
	Name() string
}

func (s *DockerfileStage) GetDependencies(_ Conveyor, _, _ image.ImageInterface) (string, error) {
	var dockerMetaArgsString []string
	for key, value := range s.dockerArgsHash {
		dockerMetaArgsString = append(dockerMetaArgsString, fmt.Sprintf("%s=%s", key, value))
	}

	shlex := shell.NewLex(parser.DefaultEscapeToken)

	var stagesDependencies [][]string
	for _, stage := range s.dockerStages {
		var dependencies []string

		dependencies = append(dependencies, s.addHost...)

		resolvedBaseName, err := shlex.ProcessWord(stage.BaseName, dockerMetaArgsString)
		if err != nil {
			return "", err
		}

		dependencies = append(dependencies, resolvedBaseName)

		for _, cmd := range stage.Commands {
			switch c := cmd.(type) {
			case *instructions.ArgCommand:
				dependencies = append(dependencies, c.String())
				if argValue, exist := s.dockerArgsHash[c.Key]; exist {
					dependencies = append(dependencies, argValue)
				}
			case *instructions.AddCommand:
				dependencies = append(dependencies, c.String())

				checksum, err := s.calculateFilesChecksum(c.SourcesAndDest.Sources())
				if err != nil {
					return "", err
				}
				dependencies = append(dependencies, checksum)
			case *instructions.CopyCommand:
				dependencies = append(dependencies, c.String())
				if c.From == "" {
					checksum, err := s.calculateFilesChecksum(c.SourcesAndDest.Sources())
					if err != nil {
						return "", err
					}
					dependencies = append(dependencies, checksum)
				}
			case dockerfileInstructionInterface:
				dependencies = append(dependencies, c.String())
			default:
				panic("runtime error")
			}
		}

		stagesDependencies = append(stagesDependencies, dependencies)
	}

	for ind, stage := range s.dockerStages {
		for relatedStageIndex, relatedStage := range s.dockerStages {
			if ind == relatedStageIndex {
				continue
			}

			if stage.BaseName == relatedStage.Name {
				stagesDependencies[ind] = append(stagesDependencies[ind], stagesDependencies[relatedStageIndex]...)
			}
		}

		for _, cmd := range stage.Commands {
			switch c := cmd.(type) {
			case *instructions.CopyCommand:
				if c.From != "" {
					relatedStageIndex, err := strconv.Atoi(c.From)
					if err != nil || relatedStageIndex >= len(stagesDependencies) {
						return "", fmt.Errorf("COPY --from refers to nonexistent dockerfile stage %s", c.From)
					} else {
						stagesDependencies[ind] = append(stagesDependencies[ind], stagesDependencies[relatedStageIndex]...)
					}
				}
			}
		}
	}

	return util.Sha256Hash(stagesDependencies[s.dockerTargetStageIndex]...), nil
}

func (s *DockerfileStage) PrepareImage(c Conveyor, prevBuiltImage, img image.ImageInterface) error {
	img.DockerfileImageBuilder().AppendBuildArgs(s.DockerBuildArgs()...)
	return nil
}

func (s *DockerfileStage) DockerBuildArgs() []string {
	var result []string

	if s.dockerfilePath != "" {
		result = append(result, fmt.Sprintf("--file=%s", s.dockerfilePath))
	}

	if s.target != "" {
		result = append(result, fmt.Sprintf("--target=%s", s.target))
	}

	if len(s.buildArgs) != 0 {
		for key, value := range s.buildArgs {
			result = append(result, fmt.Sprintf("--build-arg=%s=%v", key, value))
		}
	}

	for _, addHost := range s.addHost {
		result = append(result, fmt.Sprintf("--add-host=%s", addHost))
	}

	result = append(result, s.context)

	return result
}

func (s *DockerfileStage) calculateFilesChecksum(wildcards []string) (string, error) {
	var checksum string
	var err error

	logProcessMsg := fmt.Sprintf("Calculating files checksum (%v)", wildcards)
	logboek.Debug.LogProcessStart(logProcessMsg, logboek.LevelLogProcessStartOptions{})
	if s.localGitRepo != nil {
		checksum, err = s.calculateFilesChecksumWithLsTree(wildcards)
	} else {
		checksum, err = s.calculateFilesChecksumWithFilesRead(wildcards)
	}

	if err != nil {
		logboek.Debug.LogProcessFail(logboek.LevelLogProcessFailOptions{})
		return "", err
	}

	logboek.Debug.LogProcessEnd(logboek.LevelLogProcessEndOptions{})

	logboek.Debug.LogF("Result checksum: %s\n", checksum)
	logboek.Debug.LogOptionalLn()

	return checksum, nil
}

func (s *DockerfileStage) calculateFilesChecksumWithLsTree(wildcards []string) (string, error) {
	if s.mainLsTreeResult == nil {
		processMsg := fmt.Sprintf("ls-tree (%s)", s.dockerignorePathMatcher.String())
		logboek.Debug.LogProcessStart(processMsg, logboek.LevelLogProcessStartOptions{})
		result, err := s.localGitRepo.LsTree(s.dockerignorePathMatcher)
		if err != nil {
			if err.Error() == "entry not found" {
				logboek.Debug.LogFWithCustomStyle(
					logboek.StyleByName(logboek.FailStyleName),
					"Entry %s is not found\n",
					s.dockerignorePathMatcher.BaseFilepath(),
				)
				logboek.Debug.LogProcessEnd(logboek.LevelLogProcessEndOptions{})
				goto entryNotFoundInGitRepository
			}

			logboek.Debug.LogProcessFail(logboek.LevelLogProcessFailOptions{})
			return "", err
		}
		logboek.Debug.LogProcessEnd(logboek.LevelLogProcessEndOptions{})

		s.mainLsTreeResult = result
	}

entryNotFoundInGitRepository:
	wildcardsPathMatcher := path_matcher.NewSimplePathMatcher(s.dockerignorePathMatcher.BaseFilepath(), wildcards, false)

	var lsTreeResultChecksum string
	if s.mainLsTreeResult != nil {
		blockMsg := fmt.Sprintf("ls-tree (%s)", wildcardsPathMatcher.String())
		logboek.Debug.LogProcessStart(blockMsg, logboek.LevelLogProcessStartOptions{})
		lsTreeResult, err := s.mainLsTreeResult.LsTree(wildcardsPathMatcher)
		if err != nil {
			logboek.Debug.LogProcessFail(logboek.LevelLogProcessFailOptions{})
			return "", err
		}
		logboek.Debug.LogProcessEnd(logboek.LevelLogProcessEndOptions{})

		if !lsTreeResult.IsEmpty() {
			blockMsg = fmt.Sprintf("ls-tree result checksum (%s)", wildcardsPathMatcher.String())
			_ = logboek.Debug.LogBlock(blockMsg, logboek.LevelLogBlockOptions{}, func() error {
				lsTreeResultChecksum = lsTreeResult.Checksum()
				logboek.Debug.LogLn()
				logboek.Debug.LogLn(lsTreeResultChecksum)

				return nil
			})
		}
	}

	if s.mainStatusResult == nil {
		processMsg := fmt.Sprintf("status (%s)", s.dockerignorePathMatcher.String())
		logboek.Debug.LogProcessStart(processMsg, logboek.LevelLogProcessStartOptions{})
		result, err := s.localGitRepo.Status(s.dockerignorePathMatcher)
		if err != nil {
			logboek.Debug.LogProcessFail(logboek.LevelLogProcessFailOptions{})
			return "", err
		}
		logboek.Debug.LogProcessEnd(logboek.LevelLogProcessEndOptions{})

		s.mainStatusResult = result
	}

	blockMsg := fmt.Sprintf("status (%s)", wildcardsPathMatcher.String())
	logboek.Debug.LogProcessStart(blockMsg, logboek.LevelLogProcessStartOptions{})
	statusResult, err := s.mainStatusResult.Status(wildcardsPathMatcher)
	if err != nil {
		logboek.Debug.LogProcessFail(logboek.LevelLogProcessFailOptions{})
		return "", err
	}
	logboek.Debug.LogProcessEnd(logboek.LevelLogProcessEndOptions{})

	var statusResultChecksum string
	if !statusResult.IsEmpty() {
		blockMsg = fmt.Sprintf("Status result checksum (%s)", wildcardsPathMatcher.String())
		_ = logboek.Debug.LogBlock(blockMsg, logboek.LevelLogBlockOptions{}, func() error {
			statusResultChecksum = statusResult.Checksum()
			logboek.Debug.LogLn()
			logboek.Debug.LogLn(statusResultChecksum)
			return nil
		})
	}

	resultChecksum := util.Sha256Hash(lsTreeResultChecksum, statusResultChecksum)

	return resultChecksum, nil
}

func (s *DockerfileStage) calculateFilesChecksumWithFilesRead(wildcards []string) (string, error) {
	var dependencies []string

	for _, wildcard := range wildcards {
		contextWildcard := filepath.Join(s.context, wildcard)

		matches, err := filepath.Glob(contextWildcard)
		if err != nil {
			return "", fmt.Errorf("glob %s failed: %s", contextWildcard, err)
		}

		var fileList []string
		for _, match := range matches {
			matchFileList, err := getAllFiles(match)
			if err != nil {
				return "", fmt.Errorf("walk %s failed: %s", match, err)
			}

			fileList = append(fileList, matchFileList...)
		}

		var finalFileList []string
		for _, filePath := range fileList {
			relFilePath, err := filepath.Rel(s.projectPath, filePath)
			if err != nil {
				panic(fmt.Sprintf("unexpected condition: %s", err))
			} else if relFilePath == "." || relFilePath == ".." || strings.HasPrefix(relFilePath, ".."+string(os.PathSeparator)) {
				panic(fmt.Sprintf("unexpected condition: %s", relFilePath))
			}

			if s.dockerignorePathMatcher.MatchPath(relFilePath) {
				finalFileList = append(finalFileList, filePath)
			}
		}

		for _, filePath := range finalFileList {
			data, err := ioutil.ReadFile(filePath)
			if err != nil {
				return "", fmt.Errorf("read file %s failed: %s", filePath, err)
			}

			dependencies = append(dependencies, string(data))
			logboek.Debug.LogF("File was added: %s\n", strings.TrimPrefix(filePath, s.projectPath+string(os.PathSeparator)))
		}
	}

	checksum := util.Sha256Hash(dependencies...)

	return checksum, nil
}

func getAllFiles(target string) ([]string, error) {
	var fileList []string
	err := filepath.Walk(target, func(path string, f os.FileInfo, err error) error {
		if f.IsDir() {
			return nil
		}

		if f.Mode()&os.ModeSymlink != 0 {
			linkTo, err := os.Readlink(path)
			if err != nil {
				return err
			}

			linkFilePath := filepath.Join(filepath.Dir(path), linkTo)
			exist, err := util.FileExists(linkFilePath)
			if err != nil {
				return err
			} else if !exist {
				return nil
			} else {
				lfinfo, err := os.Stat(linkFilePath)
				if err != nil {
					return err
				}

				if lfinfo.IsDir() {
					// infinite loop detector
					if target == linkFilePath {
						return nil
					}

					lfileList, err := getAllFiles(linkFilePath)
					if err != nil {
						return err
					}

					fileList = append(fileList, lfileList...)
				} else {
					fileList = append(fileList, linkFilePath)
				}

				return nil
			}
		}

		fileList = append(fileList, path)
		return err
	})

	return fileList, err
}
<|endoftext|>


<commit_before>package aws

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"net/http"
	"time"

	"github.com/atlassian/gostatsd"

	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/spf13/viper"
	"golang.org/x/net/http2"
)

const (
	// ProviderName is the name of AWS cloud provider.
	ProviderName         = "aws"
	defaultClientTimeout = 9 * time.Second
)

// Provider represents an AWS provider.
type Provider struct {
	Metadata *ec2metadata.EC2Metadata
	Ec2      *ec2.EC2
}

func newEc2Filter(name string, value string) *ec2.Filter {
	return &ec2.Filter{
		Name: aws.String(name),
		Values: []*string{
			aws.String(value),
		},
	}
}

// Instance returns the instance details from aws.
func (p *Provider) Instance(ctx context.Context, IP gostatsd.IP) (*gostatsd.Instance, error) {
	req, _ := p.Ec2.DescribeInstancesRequest(&ec2.DescribeInstancesInput{
		Filters: []*ec2.Filter{
			newEc2Filter("private-ip-address", string(IP)),
		},
	})
	req.HTTPRequest = req.HTTPRequest.WithContext(ctx)
	var inst *ec2.Instance
	err := req.EachPage(func(data interface{}, isLastPage bool) bool {
		for _, reservation := range data.(*ec2.DescribeInstancesOutput).Reservations {
			for _, instance := range reservation.Instances {
				inst = instance
				return false
			}
		}
		return true
	})
	if err != nil {
		return nil, fmt.Errorf("error listing AWS instances: %v", err)
	}
	if inst == nil {
		return nil, errors.New("no instances found")
	}
	region, err := azToRegion(aws.StringValue(inst.Placement.AvailabilityZone))
	if err != nil {
		log.Errorf("Error getting instance region: %v", err)
	}
	tags := make(gostatsd.Tags, len(inst.Tags))
	for idx, tag := range inst.Tags {
		tags[idx] = fmt.Sprintf("%s:%s",
			gostatsd.NormalizeTagKey(aws.StringValue(tag.Key)),
			aws.StringValue(tag.Value))
	}
	instance := &gostatsd.Instance{
		ID:     aws.StringValue(inst.InstanceId),
		Region: region,
		Tags:   tags,
	}
	return instance, nil
}

// Name returns the name of the provider.
func (p *Provider) Name() string {
	return ProviderName
}

// SelfIP returns host's IPv4 address.
func (p *Provider) SelfIP() (gostatsd.IP, error) {
	ip, err := p.Metadata.GetMetadata("local-ipv4")
	return gostatsd.IP(ip), err
}

// Derives the region from a valid az name.
// Returns an error if the az is known invalid (empty).
func azToRegion(az string) (string, error) {
	if az == "" {
		return "", errors.New("invalid (empty) AZ")
	}
	region := az[:len(az)-1]
	return region, nil
}

// NewProviderFromViper returns a new aws provider.
func NewProviderFromViper(v *viper.Viper) (gostatsd.CloudProvider, error) {
	a := getSubViper(v, "aws")
	a.SetDefault("max_retries", 3)
	a.SetDefault("client_timeout", defaultClientTimeout)
	httpTimeout := a.GetDuration("client_timeout")
	if httpTimeout <= 0 {
		return nil, errors.New("client timeout must be positive")
	}

	// This is the main config without credentials.
	transport := &http.Transport{
		Proxy:               http.ProxyFromEnvironment,
		TLSHandshakeTimeout: 3 * time.Second,
		TLSClientConfig: &tls.Config{
			// Can't use SSLv3 because of POODLE and BEAST
			// Can't use TLSv1.0 because of POODLE and BEAST using CBC cipher
			// Can't use TLSv1.1 because of RC4 cipher usage
			MinVersion: tls.VersionTLS12,
		},
		DialContext: (&net.Dialer{
			Timeout:   5 * time.Second,
			KeepAlive: 30 * time.Second,
		}).DialContext,
		MaxIdleConns:    50,
		IdleConnTimeout: 1 * time.Minute,
	}
	if err := http2.ConfigureTransport(transport); err != nil {
		return nil, err
	}
	sharedConfig := aws.NewConfig().
		WithHTTPClient(&http.Client{
			Transport: transport,
			Timeout:   httpTimeout,
		}).
		WithMaxRetries(a.GetInt("max_retries"))
	metadata := ec2metadata.New(session.New(sharedConfig))
	az, err := metadata.GetMetadata("placement/availability-zone")
	if err != nil {
		return nil, fmt.Errorf("error getting availability zone: %v", err)
	}
	region, err := azToRegion(az)
	if err != nil {
		return nil, fmt.Errorf("error getting AWS region: %v", err)
	}
	ec2config := sharedConfig.Copy().
		WithCredentials(credentials.NewChainCredentials(
			[]credentials.Provider{
				&credentials.EnvProvider{},
				&ec2rolecreds.EC2RoleProvider{
					Client: metadata,
				},
				&credentials.SharedCredentialsProvider{},
			})).
		WithRegion(region)
	return &Provider{
		Metadata: metadata,
		Ec2:      ec2.New(session.New(ec2config)),
	}, nil
}

func getSubViper(v *viper.Viper, key string) *viper.Viper {
	n := v.Sub(key)
	if n == nil {
		n = viper.New()
	}
	return n
}
<commit_msg>Use Region method<commit_after>package aws

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"net/http"
	"time"

	"github.com/atlassian/gostatsd"

	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/spf13/viper"
	"golang.org/x/net/http2"
)

const (
	// ProviderName is the name of AWS cloud provider.
	ProviderName         = "aws"
	defaultClientTimeout = 9 * time.Second
)

// Provider represents an AWS provider.
type Provider struct {
	Metadata *ec2metadata.EC2Metadata
	Ec2      *ec2.EC2
}

func newEc2Filter(name string, value string) *ec2.Filter {
	return &ec2.Filter{
		Name: aws.String(name),
		Values: []*string{
			aws.String(value),
		},
	}
}

// Instance returns the instance details from aws.
func (p *Provider) Instance(ctx context.Context, IP gostatsd.IP) (*gostatsd.Instance, error) {
	req, _ := p.Ec2.DescribeInstancesRequest(&ec2.DescribeInstancesInput{
		Filters: []*ec2.Filter{
			newEc2Filter("private-ip-address", string(IP)),
		},
	})
	req.HTTPRequest = req.HTTPRequest.WithContext(ctx)
	var inst *ec2.Instance
	err := req.EachPage(func(data interface{}, isLastPage bool) bool {
		for _, reservation := range data.(*ec2.DescribeInstancesOutput).Reservations {
			for _, instance := range reservation.Instances {
				inst = instance
				return false
			}
		}
		return true
	})
	if err != nil {
		return nil, fmt.Errorf("error listing AWS instances: %v", err)
	}
	if inst == nil {
		return nil, errors.New("no instances found")
	}
	region, err := azToRegion(aws.StringValue(inst.Placement.AvailabilityZone))
	if err != nil {
		log.Errorf("Error getting instance region: %v", err)
	}
	tags := make(gostatsd.Tags, len(inst.Tags))
	for idx, tag := range inst.Tags {
		tags[idx] = fmt.Sprintf("%s:%s",
			gostatsd.NormalizeTagKey(aws.StringValue(tag.Key)),
			aws.StringValue(tag.Value))
	}
	instance := &gostatsd.Instance{
		ID:     aws.StringValue(inst.InstanceId),
		Region: region,
		Tags:   tags,
	}
	return instance, nil
}

// Name returns the name of the provider.
func (p *Provider) Name() string {
	return ProviderName
}

// SelfIP returns host's IPv4 address.
func (p *Provider) SelfIP() (gostatsd.IP, error) {
	ip, err := p.Metadata.GetMetadata("local-ipv4")
	return gostatsd.IP(ip), err
}

// Derives the region from a valid az name.
// Returns an error if the az is known invalid (empty).
func azToRegion(az string) (string, error) {
	if az == "" {
		return "", errors.New("invalid (empty) AZ")
	}
	region := az[:len(az)-1]
	return region, nil
}

// NewProviderFromViper returns a new aws provider.
func NewProviderFromViper(v *viper.Viper) (gostatsd.CloudProvider, error) {
	a := getSubViper(v, "aws")
	a.SetDefault("max_retries", 3)
	a.SetDefault("client_timeout", defaultClientTimeout)
	httpTimeout := a.GetDuration("client_timeout")
	if httpTimeout <= 0 {
		return nil, errors.New("client timeout must be positive")
	}

	// This is the main config without credentials.
	transport := &http.Transport{
		Proxy:               http.ProxyFromEnvironment,
		TLSHandshakeTimeout: 3 * time.Second,
		TLSClientConfig: &tls.Config{
			// Can't use SSLv3 because of POODLE and BEAST
			// Can't use TLSv1.0 because of POODLE and BEAST using CBC cipher
			// Can't use TLSv1.1 because of RC4 cipher usage
			MinVersion: tls.VersionTLS12,
		},
		DialContext: (&net.Dialer{
			Timeout:   5 * time.Second,
			KeepAlive: 30 * time.Second,
		}).DialContext,
		MaxIdleConns:    50,
		IdleConnTimeout: 1 * time.Minute,
	}
	if err := http2.ConfigureTransport(transport); err != nil {
		return nil, err
	}
	sharedConfig := aws.NewConfig().
		WithHTTPClient(&http.Client{
			Transport: transport,
			Timeout:   httpTimeout,
		}).
		WithMaxRetries(a.GetInt("max_retries"))
	metadata := ec2metadata.New(session.New(sharedConfig))
	region, err := metadata.Region()
	if err != nil {
		return nil, fmt.Errorf("error getting AWS region: %v", err)
	}
	ec2config := sharedConfig.Copy().
		WithCredentials(credentials.NewChainCredentials(
			[]credentials.Provider{
				&credentials.EnvProvider{},
				&ec2rolecreds.EC2RoleProvider{
					Client: metadata,
				},
				&credentials.SharedCredentialsProvider{},
			})).
		WithRegion(region)
	return &Provider{
		Metadata: metadata,
		Ec2:      ec2.New(session.New(ec2config)),
	}, nil
}

func getSubViper(v *viper.Viper, key string) *viper.Viper {
	n := v.Sub(key)
	if n == nil {
		n = viper.New()
	}
	return n
}
<|endoftext|>


<commit_before>// Copyright 2019 Google LLC 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 options

import (
	"runtime"

	"github.com/spf13/cobra"
)

// BuildOptions represents options for the ko builder.
type BuildOptions struct {
	ConcurrentBuilds     int
	DisableOptimizations bool
}

func AddBuildOptions(cmd *cobra.Command, bo *BuildOptions) {
	cmd.Flags().IntVarP(&bo.ConcurrentBuilds, "--jobs", "j", runtime.GOMAXPROCS(0),
		"The maximum number of concurrent builds")
	cmd.Flags().BoolVar(&bo.DisableOptimizations, "disable-optimizations", bo.DisableOptimizations,
		"Disable optimizations when building Go code. Useful when you want to interactively debug the created container.")
}
<commit_msg>Remove extra dashes (#86)<commit_after>// Copyright 2019 Google LLC 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 options

import (
	"runtime"

	"github.com/spf13/cobra"
)

// BuildOptions represents options for the ko builder.
type BuildOptions struct {
	ConcurrentBuilds     int
	DisableOptimizations bool
}

func AddBuildOptions(cmd *cobra.Command, bo *BuildOptions) {
	cmd.Flags().IntVarP(&bo.ConcurrentBuilds, "jobs", "j", runtime.GOMAXPROCS(0),
		"The maximum number of concurrent builds")
	cmd.Flags().BoolVar(&bo.DisableOptimizations, "disable-optimizations", bo.DisableOptimizations,
		"Disable optimizations when building Go code. Useful when you want to interactively debug the created container.")
}
<|endoftext|>


<commit_before>package httpclient

import (
	"crypto/tls"
	"crypto/x509"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestGet(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, "GET", r.Method)
		assert.Equal(t, "/path", r.URL.Path)
		w.Write([]byte("ok"))
	}))
	defer ts.Close()

	client := New(ts.URL)

	resp, err := client.Get("/path")
	require.NoError(t, err)

	respBody, err := ioutil.ReadAll(resp.Body)
	require.NoError(t, err)
	require.Equal(t, "ok", string(respBody))
}

func TestPost(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, "POST", r.Method)
		assert.Equal(t, "/path", r.URL.Path)
		assert.Equal(t, "application/json", r.Header.Get("Content-Type"))

		body, err := ioutil.ReadAll(r.Body)
		assert.NoError(t, err)

		assert.Equal(t, `{"cluster":"DC/OS"}`, string(body))
		w.Write([]byte("ok"))
	}))
	defer ts.Close()

	client := New(ts.URL)

	resp, err := client.Post("/path", "application/json", strings.NewReader(`{"cluster":"DC/OS"}`))
	require.NoError(t, err)

	respBody, err := ioutil.ReadAll(resp.Body)
	require.NoError(t, err)
	require.Equal(t, "ok", string(respBody))
}

func TestNewRequest(t *testing.T) {
	client := New("https://dcos.io", func(client *Client) {
		client.acsToken = "acsToken"
	})

	req, err := client.NewRequest("GET", "/path", nil)
	require.NoError(t, err)
	require.Equal(t, req.URL.String(), "https://dcos.io/path")
	require.Equal(t, "token=acsToken", req.Header.Get("Authorization"))
}

func TestTimeout(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		timeout, err := time.ParseDuration(r.URL.Query().Get("timeout"))
		assert.NoError(t, err)
		time.Sleep(timeout)
	}))
	defer ts.Close()

	client := New(ts.URL, func(client *Client) {
		client.timeout = 50 * time.Millisecond
	})

	// The handler will sleep for 100ms with a client timeout of 50ms, the call should fail.
	req, err := client.NewRequest("GET", "/", nil)
	req.URL.RawQuery = "timeout=100ms"
	require.NoError(t, err)
	_, err = client.Do(req)
	require.Error(t, err)

	// The handler will sleep for 10ms with a client timeout of 50ms, the call should succeed.
	req, err = client.NewRequest("GET", "/", nil)
	req.URL.RawQuery = "timeout=10ms"
	require.NoError(t, err)
	_, err = client.Do(req)
	require.NoError(t, err)
}

func TestTLS(t *testing.T) {
	ts := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("ok"))
	}))
	defer ts.Close()

	certPool := x509.NewCertPool()
	certPool.AddCert(ts.Certificate())

	tlsConfigs := []struct {
		tls   *tls.Config
		valid bool
	}{
		// Using an empty TLS config should fail because the CA is not specified.
		{&tls.Config{}, false},

		// Using a TLS config with the actual CA should work.
		{&tls.Config{RootCAs: certPool}, true},

		// Using a TLS config with InsecureSkipVerify set to true should work.
		{&tls.Config{InsecureSkipVerify: true}, true},
	}

	for _, exp := range tlsConfigs {
		client := New(ts.URL, func(client *Client) {
			client.baseClient.Transport = &http.Transport{
				TLSClientConfig: exp.tls,
			}
		})

		resp, err := client.Get("/")
		if exp.valid {
			require.NoError(t, err)
			respBody, err := ioutil.ReadAll(resp.Body)
			require.NoError(t, err)
			require.Equal(t, "ok", string(respBody))
		} else {
			require.Error(t, err)
			require.Nil(t, resp)
		}
	}
}
<commit_msg>httpclient: update TestTimeout not to be network-dependent<commit_after>package httpclient

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestGet(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, "GET", r.Method)
		assert.Equal(t, "/path", r.URL.Path)
		w.Write([]byte("ok"))
	}))
	defer ts.Close()

	client := New(ts.URL)

	resp, err := client.Get("/path")
	require.NoError(t, err)

	respBody, err := ioutil.ReadAll(resp.Body)
	require.NoError(t, err)
	require.Equal(t, "ok", string(respBody))
}

func TestPost(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, "POST", r.Method)
		assert.Equal(t, "/path", r.URL.Path)
		assert.Equal(t, "application/json", r.Header.Get("Content-Type"))

		body, err := ioutil.ReadAll(r.Body)
		assert.NoError(t, err)

		assert.Equal(t, `{"cluster":"DC/OS"}`, string(body))
		w.Write([]byte("ok"))
	}))
	defer ts.Close()

	client := New(ts.URL)

	resp, err := client.Post("/path", "application/json", strings.NewReader(`{"cluster":"DC/OS"}`))
	require.NoError(t, err)

	respBody, err := ioutil.ReadAll(resp.Body)
	require.NoError(t, err)
	require.Equal(t, "ok", string(respBody))
}

func TestNewRequest(t *testing.T) {
	client := New("https://dcos.io", func(client *Client) {
		client.acsToken = "acsToken"
	})

	req, err := client.NewRequest("GET", "/path", nil)
	require.NoError(t, err)
	require.Equal(t, req.URL.String(), "https://dcos.io/path")
	require.Equal(t, "token=acsToken", req.Header.Get("Authorization"))
}

func TestCancelRequest(t *testing.T) {
	done := make(chan struct{})
	stuckHandler := make(chan struct{})
	canceler := make(chan context.CancelFunc, 1)

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		cancel := <-canceler
		cancel()
		<-stuckHandler
	}))
	defer ts.Close()
	defer close(stuckHandler)

	client := New(ts.URL)

	req, err := client.NewRequest("GET", "/", nil)
	require.NoError(t, err)

	// Create a cancelable request and send the cancel function to a channel.
	// The HTTP handler will then invoke it, this simulates a test where the
	// request timeout is reached while the server is still processing it.
	newCtx, cancel := context.WithCancel(req.Context())
	req = req.WithContext(newCtx)
	canceler <- cancel

	go func() {
		resp, err := client.Do(req)
		require.Error(t, err)
		require.Nil(t, resp)
		close(done)
	}()

	select {
	case <-time.After(5 * time.Second):
		require.Fail(t, "HTTP client didn't error-out within 5 seconds, it is most likely stuck forever.")
	case <-done:
	}
}

func TestTLS(t *testing.T) {
	ts := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("ok"))
	}))
	defer ts.Close()

	certPool := x509.NewCertPool()
	certPool.AddCert(ts.Certificate())

	tlsConfigs := []struct {
		tls   *tls.Config
		valid bool
	}{
		// Using an empty TLS config should fail because the CA is not specified.
		{&tls.Config{}, false},

		// Using a TLS config with the actual CA should work.
		{&tls.Config{RootCAs: certPool}, true},

		// Using a TLS config with InsecureSkipVerify set to true should work.
		{&tls.Config{InsecureSkipVerify: true}, true},
	}

	for _, exp := range tlsConfigs {
		client := New(ts.URL, func(client *Client) {
			client.baseClient.Transport = &http.Transport{
				TLSClientConfig: exp.tls,
			}
		})

		resp, err := client.Get("/")
		if exp.valid {
			require.NoError(t, err)
			respBody, err := ioutil.ReadAll(resp.Body)
			require.NoError(t, err)
			require.Equal(t, "ok", string(respBody))
		} else {
			require.Error(t, err)
			require.Nil(t, resp)
		}
	}
}
<|endoftext|>


<commit_before>package client

import (
	"flag"
	"io"
	"time"

	"github.com/cortexproject/cortex/pkg/distributor"
	"github.com/cortexproject/cortex/pkg/util/grpcclient"
	"github.com/grpc-ecosystem/grpc-opentracing/go/otgrpc"
	"github.com/opentracing/opentracing-go"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/weaveworks/common/middleware"
	"google.golang.org/grpc"
	"google.golang.org/grpc/health/grpc_health_v1"

	cortex_middleware "github.com/cortexproject/cortex/pkg/util/middleware"

	"github.com/grafana/loki/pkg/logproto"
)

var ingesterClientRequestDuration = promauto.NewHistogramVec(prometheus.HistogramOpts{
	Name:    "loki_ingester_client_request_duration_seconds",
	Help:    "Time spent doing Ingester requests.",
	Buckets: prometheus.ExponentialBuckets(0.001, 4, 6),
}, []string{"operation", "status_code"})

type HealthAndIngesterClient interface {
	logproto.IngesterClient
	grpc_health_v1.HealthClient
	Close() error
}

type ClosableHealthAndIngesterClient struct {
	logproto.PusherClient
	logproto.QuerierClient
	logproto.IngesterClient
	grpc_health_v1.HealthClient
	io.Closer
}

// Config for an ingester client.
type Config struct {
	PoolConfig                   distributor.PoolConfig         `yaml:"pool_config,omitempty"`
	RemoteTimeout                time.Duration                  `yaml:"remote_timeout,omitempty"`
	GRPCClientConfig             grpcclient.Config              `yaml:"grpc_client_config"`
	GRPCUnaryClientInterceptors  []grpc.UnaryClientInterceptor  `yaml:"-"`
	GRCPStreamClientInterceptors []grpc.StreamClientInterceptor `yaml:"-"`
}

// RegisterFlags registers flags.
func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
	cfg.GRPCClientConfig.RegisterFlagsWithPrefix("ingester.client", f)
	cfg.PoolConfig.RegisterFlags(f)

	f.DurationVar(&cfg.PoolConfig.RemoteTimeout, "ingester.client.healthcheck-timeout", 1*time.Second, "Timeout for healthcheck rpcs.")
	f.DurationVar(&cfg.RemoteTimeout, "ingester.client.timeout", 5*time.Second, "Timeout for ingester client RPCs.")
}

// New returns a new ingester client.
func New(cfg Config, addr string) (HealthAndIngesterClient, error) {
	opts := []grpc.DialOption{
		grpc.WithInsecure(),
		grpc.WithDefaultCallOptions(cfg.GRPCClientConfig.CallOptions()...),
	}

	dialOpts, err := cfg.GRPCClientConfig.DialOption(instrumentation(&cfg))
	if err != nil {
		return nil, err
	}

	opts = append(opts, dialOpts...)
	conn, err := grpc.Dial(addr, opts...)
	if err != nil {
		return nil, err
	}
	return ClosableHealthAndIngesterClient{
		PusherClient:   logproto.NewPusherClient(conn),
		QuerierClient:  logproto.NewQuerierClient(conn),
		IngesterClient: logproto.NewIngesterClient(conn),
		HealthClient:   grpc_health_v1.NewHealthClient(conn),
		Closer:         conn,
	}, nil
}

func instrumentation(cfg *Config) ([]grpc.UnaryClientInterceptor, []grpc.StreamClientInterceptor) {
	var unaryInterceptors []grpc.UnaryClientInterceptor
	unaryInterceptors = append(unaryInterceptors, cfg.GRPCUnaryClientInterceptors...)
	unaryInterceptors = append(unaryInterceptors,
		otgrpc.OpenTracingClientInterceptor(opentracing.GlobalTracer()),
		middleware.ClientUserHeaderInterceptor,
		cortex_middleware.PrometheusGRPCUnaryInstrumentation(ingesterClientRequestDuration),
	)
	var streamInterceptors []grpc.StreamClientInterceptor
	streamInterceptors = append(streamInterceptors, cfg.GRCPStreamClientInterceptors...)
	streamInterceptors = append(streamInterceptors,
		otgrpc.OpenTracingStreamClientInterceptor(opentracing.GlobalTracer()),
		middleware.StreamClientUserHeaderInterceptor,
		cortex_middleware.PrometheusGRPCStreamInstrumentation(ingesterClientRequestDuration),
	)

	return unaryInterceptors, streamInterceptors
}
<commit_msg>Authc/z: Enable grpc_client_config to allow mTLS (#4176)<commit_after>package client

import (
	"flag"
	"io"
	"time"

	"github.com/cortexproject/cortex/pkg/distributor"
	"github.com/cortexproject/cortex/pkg/util/grpcclient"
	"github.com/grpc-ecosystem/grpc-opentracing/go/otgrpc"
	"github.com/opentracing/opentracing-go"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/weaveworks/common/middleware"
	"google.golang.org/grpc"
	"google.golang.org/grpc/health/grpc_health_v1"

	cortex_middleware "github.com/cortexproject/cortex/pkg/util/middleware"

	"github.com/grafana/loki/pkg/logproto"
)

var ingesterClientRequestDuration = promauto.NewHistogramVec(prometheus.HistogramOpts{
	Name:    "loki_ingester_client_request_duration_seconds",
	Help:    "Time spent doing Ingester requests.",
	Buckets: prometheus.ExponentialBuckets(0.001, 4, 6),
}, []string{"operation", "status_code"})

type HealthAndIngesterClient interface {
	logproto.IngesterClient
	grpc_health_v1.HealthClient
	Close() error
}

type ClosableHealthAndIngesterClient struct {
	logproto.PusherClient
	logproto.QuerierClient
	logproto.IngesterClient
	grpc_health_v1.HealthClient
	io.Closer
}

// Config for an ingester client.
type Config struct {
	PoolConfig                   distributor.PoolConfig         `yaml:"pool_config,omitempty"`
	RemoteTimeout                time.Duration                  `yaml:"remote_timeout,omitempty"`
	GRPCClientConfig             grpcclient.Config              `yaml:"grpc_client_config"`
	GRPCUnaryClientInterceptors  []grpc.UnaryClientInterceptor  `yaml:"-"`
	GRCPStreamClientInterceptors []grpc.StreamClientInterceptor `yaml:"-"`
}

// RegisterFlags registers flags.
func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
	cfg.GRPCClientConfig.RegisterFlagsWithPrefix("ingester.client", f)
	cfg.PoolConfig.RegisterFlags(f)

	f.DurationVar(&cfg.PoolConfig.RemoteTimeout, "ingester.client.healthcheck-timeout", 1*time.Second, "Timeout for healthcheck rpcs.")
	f.DurationVar(&cfg.RemoteTimeout, "ingester.client.timeout", 5*time.Second, "Timeout for ingester client RPCs.")
}

// New returns a new ingester client.
func New(cfg Config, addr string) (HealthAndIngesterClient, error) {
	opts := []grpc.DialOption{
		grpc.WithDefaultCallOptions(cfg.GRPCClientConfig.CallOptions()...),
	}

	dialOpts, err := cfg.GRPCClientConfig.DialOption(instrumentation(&cfg))
	if err != nil {
		return nil, err
	}

	opts = append(opts, dialOpts...)
	conn, err := grpc.Dial(addr, opts...)
	if err != nil {
		return nil, err
	}
	return ClosableHealthAndIngesterClient{
		PusherClient:   logproto.NewPusherClient(conn),
		QuerierClient:  logproto.NewQuerierClient(conn),
		IngesterClient: logproto.NewIngesterClient(conn),
		HealthClient:   grpc_health_v1.NewHealthClient(conn),
		Closer:         conn,
	}, nil
}

func instrumentation(cfg *Config) ([]grpc.UnaryClientInterceptor, []grpc.StreamClientInterceptor) {
	var unaryInterceptors []grpc.UnaryClientInterceptor
	unaryInterceptors = append(unaryInterceptors, cfg.GRPCUnaryClientInterceptors...)
	unaryInterceptors = append(unaryInterceptors,
		otgrpc.OpenTracingClientInterceptor(opentracing.GlobalTracer()),
		middleware.ClientUserHeaderInterceptor,
		cortex_middleware.PrometheusGRPCUnaryInstrumentation(ingesterClientRequestDuration),
	)
	var streamInterceptors []grpc.StreamClientInterceptor
	streamInterceptors = append(streamInterceptors, cfg.GRCPStreamClientInterceptors...)
	streamInterceptors = append(streamInterceptors,
		otgrpc.OpenTracingStreamClientInterceptor(opentracing.GlobalTracer()),
		middleware.StreamClientUserHeaderInterceptor,
		cortex_middleware.PrometheusGRPCStreamInstrumentation(ingesterClientRequestDuration),
	)

	return unaryInterceptors, streamInterceptors
}
<|endoftext|>


<commit_before>package manifest

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestGetManifestWithoutInterval(t *testing.T) {
	handler := func(w http.ResponseWriter, r *http.Request) {
		data := `{
			"targets": [
				{
					"url": "http://www.canary.io",
					"name": "canary"
				}
			]
		}`

		fmt.Fprintf(w, data)
	}

	ts := httptest.NewServer(http.HandlerFunc(handler))
	defer ts.Close()

	m, err := GetManifest(ts.URL)
	if err != nil {
		t.Fatal(err)
	}

	if len(m.Targets) != 1 {
		t.Fatal("%d targets found, but expected 1", len(m.Targets))
	}

	target := m.Targets[0]
	if target.Name != "canary" {
		t.Fatal("expected name to be equal to 'canary', go %s", target.URL)
	}

	if target.URL != "http://www.canary.io" {
		t.Fatal("expected URL to be equal to 'http://www.canary.io', got %s", target.URL)
	}

	if target.Interval != 0 {
		t.Fatal("expected Interval to be equal to zero when undefined in the manifest json, got %d", target.Interval)
	}
}

func TestGetManifestWithInterval(t *testing.T) {
	handler := func(w http.ResponseWriter, r *http.Request) {
		data := `{
			"targets": [
				{
					"url": "http://www.canary.io",
					"name": "canary",
					"interval": 2
				},
				{
					"url": "http://www.github.com",
					"name": "github",
					"interval": 4
				}
			]
		}`

		fmt.Fprintf(w, data)
	}

	ts := httptest.NewServer(http.HandlerFunc(handler))
	defer ts.Close()

	m, err := GetManifest(ts.URL)
	if err != nil {
		t.Fatal(err)
	}

	first_target := m.Targets[0]

	if first_target.Interval != 2 {
		t.Fatal("expected Interval on the first target to be equal to the manifest json definition of 2, got %d", first_target.Interval)
	}

	second_target := m.Targets[1]

	if second_target.Interval != 4 {
		t.Fatal("expected Interval on the second target to be equal to the manifest json definition of 4, got %d", second_target.Interval)
	}
}

func TestGetManifestRampup(t *testing.T) {
	handler := func(w http.ResponseWriter, r *http.Request) {
		data := `{
			"targets": [
				{
					"url": "http://www.canary.io",
					"name": "canary"
				},
				{
					"url": "http://www.github.com",
					"name": "github"
				},
				{
					"url": "http://www.google.com",
					"name": "google"
				},
				{
					"url": "http://www.youtube.com",
					"name": "youtube"
				}
			]
		}`

		fmt.Fprintf(w, data)
	}

	ts := httptest.NewServer(http.HandlerFunc(handler))
	defer ts.Close()

	m, err := GetManifest(ts.URL)
	if err != nil {
		t.Fatal(err)
	}

	// m.StartDelays will the same length as m.Targets
	if len(m.Targets) != len(m.StartDelays) {
		t.Fatal("expected length of m.StartDelays (%d) to match length of m.Targets (%d)", len(m.Targets), len(m.StartDelays))
	}

	// without calling GenerateRampupDelays, StartDelays should all be zero.
	for index, value := range m.StartDelays {
		if value != 0.0 {
			t.Fatal("Expected initial start delay to be 0.0, got %d for index %d", value, index)
		}
	}

	// Calling GenerateRampupDelays will update m.StartDelays to an even distribution of
	// millisecond delays for the passed Sampling interval (in seconds).
	m.GenerateRampupDelays(10)

	if m.StartDelays[0] != 0.0 {
		t.Fatal("The first start delay should be 0.0 even after generation, got %d", m.StartDelays[0])
	}

	if m.StartDelays[1] != 2500.0 {
		t.Fatal("The second start delay should be 2500.0 ms after generation, got %d", m.StartDelays[1])
	}

	if m.StartDelays[2] != 5000.0 {
		t.Fatal("The second start delay should be 5000.0 ms after generation, got %d", m.StartDelays[2])
	}

	if m.StartDelays[3] != 7500.0 {
		t.Fatal("The second start delay should be 7500.0 ms after generation, got %d", m.StartDelays[3])
	}
}
<commit_msg>Fix formatting of test failure messages<commit_after>package manifest

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestGetManifestWithoutInterval(t *testing.T) {
	handler := func(w http.ResponseWriter, r *http.Request) {
		data := `{
			"targets": [
				{
					"url": "http://www.canary.io",
					"name": "canary"
				}
			]
		}`

		fmt.Fprintf(w, data)
	}

	ts := httptest.NewServer(http.HandlerFunc(handler))
	defer ts.Close()

	m, err := GetManifest(ts.URL)
	if err != nil {
		t.Fatal(err)
	}

	if len(m.Targets) != 1 {
		t.Fatalf("%d targets found, but expected 1", len(m.Targets))
	}

	target := m.Targets[0]
	if target.Name != "canary" {
		t.Fatalf("expected name to be equal to 'canary', go %s", target.URL)
	}

	if target.URL != "http://www.canary.io" {
		t.Fatalf("expected URL to be equal to 'http://www.canary.io', got %s", target.URL)
	}

	if target.Interval != 0 {
		t.Fatalf("expected Interval to be equal to zero when undefined in the manifest json, got %d", target.Interval)
	}
}

func TestGetManifestWithInterval(t *testing.T) {
	handler := func(w http.ResponseWriter, r *http.Request) {
		data := `{
			"targets": [
				{
					"url": "http://www.canary.io",
					"name": "canary",
					"interval": 2
				},
				{
					"url": "http://www.github.com",
					"name": "github",
					"interval": 4
				}
			]
		}`

		fmt.Fprintf(w, data)
	}

	ts := httptest.NewServer(http.HandlerFunc(handler))
	defer ts.Close()

	m, err := GetManifest(ts.URL)
	if err != nil {
		t.Fatal(err)
	}

	first_target := m.Targets[0]

	if first_target.Interval != 2 {
		t.Fatalf("expected Interval on the first target to be equal to the manifest json definition of 2, got %d", first_target.Interval)
	}

	second_target := m.Targets[1]

	if second_target.Interval != 4 {
		t.Fatalf("expected Interval on the second target to be equal to the manifest json definition of 4, got %d", second_target.Interval)
	}
}

func TestGetManifestRampup(t *testing.T) {
	handler := func(w http.ResponseWriter, r *http.Request) {
		data := `{
			"targets": [
				{
					"url": "http://www.canary.io",
					"name": "canary"
				},
				{
					"url": "http://www.github.com",
					"name": "github"
				},
				{
					"url": "http://www.google.com",
					"name": "google"
				},
				{
					"url": "http://www.youtube.com",
					"name": "youtube"
				}
			]
		}`

		fmt.Fprintf(w, data)
	}

	ts := httptest.NewServer(http.HandlerFunc(handler))
	defer ts.Close()

	m, err := GetManifest(ts.URL)
	if err != nil {
		t.Fatal(err)
	}

	// m.StartDelays will the same length as m.Targets
	if len(m.Targets) != len(m.StartDelays) {
		t.Fatalf("expected length of m.StartDelays (%d) to match length of m.Targets (%d)", len(m.Targets), len(m.StartDelays))
	}

	// without calling GenerateRampupDelays, StartDelays should all be zero.
	for index, value := range m.StartDelays {
		if value != 0.0 {
			t.Fatalf("Expected initial start delay to be 0.0, got %d for index %d", value, index)
		}
	}

	// Calling GenerateRampupDelays will update m.StartDelays to an even distribution of
	// millisecond delays for the passed Sampling interval (in seconds).
	m.GenerateRampupDelays(10)

	if m.StartDelays[0] != 0.0 {
		t.Fatalf("The first start delay should be 0.0 even after generation, got %d", m.StartDelays[0])
	}

	if m.StartDelays[1] != 2500.0 {
		t.Fatalf("The second start delay should be 2500.0 ms after generation, got %d", m.StartDelays[1])
	}

	if m.StartDelays[2] != 5000.0 {
		t.Fatalf("The second start delay should be 5000.0 ms after generation, got %d", m.StartDelays[2])
	}

	if m.StartDelays[3] != 7500.0 {
		t.Fatalf("The second start delay should be 7500.0 ms after generation, got %d", m.StartDelays[3])
	}
}
<|endoftext|>