<commit_before>package cache

import (
	"errors"
	"fmt"
	"math/rand"
	"sort"
	"testing"
	"time"

	"github.com/grafana/metrictank/mdata/cache/accnt"
	"github.com/grafana/metrictank/mdata/chunk"
	"github.com/grafana/metrictank/test"
	"github.com/raintank/schema"
)

func generateChunks(b testing.TB, startAt, count, step uint32) []chunk.IterGen {
	res := make([]chunk.IterGen, 0, count)

	values := make([]uint32, step)
	for t0 := startAt; t0 < startAt+(step*uint32(count)); t0 += step {
		c := getItgen(b, values, t0, true)
		res = append(res, c)
	}
	return res
}

func getCCM() (schema.AMKey, *CCacheMetric) {
	amkey, _ := schema.AMKeyFromString("1.12345678901234567890123456789012")
	ccm := NewCCacheMetric(amkey.MKey)
	return amkey, ccm
}

// TestAddAsc tests adding ascending timestamp chunks individually
func TestAddAsc(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		prev := uint32(1)
		for _, chunk := range chunks {
			ccm.Add(prev, chunk)
			prev = chunk.Ts
		}
	})
}

// TestAddDesc1 tests adding chunks that are all descending
func TestAddDesc1(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		for i := len(chunks) - 1; i >= 0; i-- {
			ccm.Add(0, chunks[i])
		}
	})
}

// TestAddDesc4 tests adding chunks that are globally descending
// but in groups 4 are ascending
func TestAddDesc4(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.Add(0, chunks[2])
		ccm.Add(0, chunks[3])
		ccm.Add(0, chunks[4])
		ccm.Add(0, chunks[5])
		ccm.Add(0, chunks[0])
		ccm.Add(0, chunks[1])
	})
}

// TestAddRange tests adding a contiguous range at once
func TestAddRange(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		prev := uint32(10)
		ccm.AddRange(prev, chunks)
	})
}

// TestAddRangeDesc4 benchmarks adding chunks that are globally descending
// but in groups of 4 are ascending. those groups are added via 1 AddRange.
func TestAddRangeDesc4(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.AddRange(0, chunks[2:6])
		ccm.AddRange(0, chunks[0:2])
	})
}

// test executes the run function
// run should generate chunks and add them to the CCacheMetric however it likes,
// but in a way so that the result will be as expected
func testRun(t *testing.T, run func(*CCacheMetric)) {
	amkey, ccm := getCCM()

	run(ccm)

	res := CCSearchResult{}
	ccm.Search(test.NewContext(), amkey, &res, 25, 45)
	if res.Complete != true {
		t.Fatalf("Expected result to be complete, but it was not")
	}

	if res.Start[0].Ts != 20 {
		t.Fatalf("Expected result to start at 20, but had %d", res.Start[0].Ts)
	}

	if res.Start[len(res.Start)-1].Ts != 40 {
		t.Fatalf("Expected result to start at 40, but had %d", res.Start[len(res.Start)-1].Ts)
	}
}

// BenchmarkAddAsc benchmarks adding ascending timestamp chunks individually
func BenchmarkAddAsc(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	prev := uint32(1)
	b.ResetTimer()
	for _, chunk := range chunks {
		ccm.Add(prev, chunk)
		prev = chunk.Ts
	}
}

// BenchmarkAddDesc1 benchmarks adding chunks that are all descending
func BenchmarkAddDesc1(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 1; i >= 0; i-- {
		ccm.Add(0, chunks[i])
	}
}

// BenchmarkAddDesc4 benchmarks adding chunks that are globally descending
// but in groups 4 are ascending
func BenchmarkAddDesc4(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 4)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 4; i >= 0; i -= 4 {
		ccm.Add(0, chunks[i])
		ccm.Add(0, chunks[i+1])
		ccm.Add(0, chunks[i+2])
		ccm.Add(0, chunks[i+3])
	}
}

// BenchmarkAddDesc64 benchmarks adding chunks that are globally descending
// but in groups 64 are ascending
func BenchmarkAddDesc64(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 64)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 64; i >= 0; i -= 64 {
		for offset := 0; offset < 64; offset += 1 {
			ccm.Add(0, chunks[i+offset])
		}
	}

}

// BenchmarkAddRangeAsc benchmarks adding a contiguous range at once
func BenchmarkAddRangeAsc(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	prev := uint32(1)
	b.ResetTimer()
	ccm.AddRange(prev, chunks)
}

// BenchmarkAddRangeDesc4 benchmarks adding chunks that are globally descending
// but in groups 4 are ascending. those groups are added via 1 AddRange.
func BenchmarkAddRangeDesc4(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 4)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 4; i >= 0; i -= 4 {
		ccm.AddRange(0, chunks[i:i+4])
	}
}

// BenchmarkAddRangeDesc64 benchmarks adding chunks that are globally descending
// but in groups 64 are ascending. those groups are added via 1 AddRange.
func BenchmarkAddRangeDesc64(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 64)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 64; i >= 0; i -= 64 {
		ccm.AddRange(0, chunks[i:i+64])
	}
}

func TestCorruptionCase1(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.AddRange(0, chunks[3:6])
		ccm.AddRange(0, chunks[0:4])
		if err := verifyCcm(ccm); err != nil {
			t.Fatal(err)
		}
	})
}

func getRandomNumber(min, max int) int {
	return rand.Intn(max-min) + min
}

// getRandomRange returns a range start-end so that
// end >= start and both numbers drawn from [min, max)
func getRandomRange(min, max int) (int, int) {
	number1 := getRandomNumber(min, max)
	number2 := getRandomNumber(min, max)
	if number1 > number2 {
		return number2, number1
	} else {
		return number1, number2
	}
}

func TestCorruptionCase2(t *testing.T) {
	rand.Seed(time.Now().Unix())
	_, ccm := getCCM()
	iterations := 100000

	// 100 chunks, first t0=t10, last is t0=1000
	chunks := generateChunks(t, 10, 100, 10)

	var opAdd, opAddRange, opDel, opDelRange int
	var adds, dels int

	for i := 0; i < iterations; i++ {
		// 0 = Add
		// 1 = AddRange
		// 2 = Del
		// 3 = Del range (via multi del cals)
		action := getRandomNumber(0, 4)
		switch action {
		case 0:
			chunk := getRandomNumber(0, 100)
			t.Logf("adding chunk %d", chunk)
			ccm.Add(0, chunks[chunk])
			opAdd++
			adds++
		case 1:
			from, to := getRandomRange(0, 100)
			t.Logf("adding range %d-%d", from, to)
			ccm.AddRange(0, chunks[from:to])
			adds += (to - from)
			opAddRange++
		case 2:
			chunk := getRandomNumber(0, 100)
			t.Logf("deleting chunk %d", chunk)
			ccm.Del(chunks[chunk].Ts) // note: chunk may not exist
			opDel++
			dels++
		case 3:
			from, to := getRandomRange(0, 100)
			t.Logf("deleting range %d-%d", from, to)
			for chunk := from; chunk < to; chunk++ {
				ccm.Del(chunks[chunk].Ts) // note: chunk may not exist
			}
			opDelRange++
			dels += (to - from)
		}

		if err := verifyCcm(ccm); err != nil {
			t.Fatal(err)
		}
	}

	fmt.Printf("operations: add %d - addRange %d - del %d - delRange %d\n", opAdd, opAddRange, opDel, opDelRange)
	fmt.Printf("total chunk adds %d - total chunk deletes %d\n", adds, dels)
}

// verifyCcm verifies the integrity of a CCacheMetric
// it assumes that all itergens are span-aware
func verifyCcm(ccm *CCacheMetric) error {
	var chunk *CCacheChunk
	var ok bool

	if len(ccm.chunks) != len(ccm.keys) {
		return errors.New("Length of ccm.chunks does not match ccm.keys")
	}

	if !sort.IsSorted(accnt.Uint32Asc(ccm.keys)) {
		return errors.New("keys are not sorted")
	}

	for i, ts := range ccm.keys {
		if chunk, ok = ccm.chunks[ts]; !ok {
			return fmt.Errorf("Ts %d is in ccm.keys but not in ccm.chunks", ts)
		}

		if i == 0 {
			if chunk.Prev != 0 {
				return errors.New("First chunk has Prev != 0")
			}
		} else {
			if chunk.Prev == 0 {
				if ccm.chunks[ccm.keys[i-1]].Ts == chunk.Ts-chunk.Itgen.Span {
					return fmt.Errorf("Chunk of ts %d has Prev == 0, but the previous chunk is present", ts)
				}
			} else {
				if ccm.chunks[ccm.keys[i-1]].Ts != chunk.Prev {
					return fmt.Errorf("Chunk of ts %d has Prev set to wrong ts %d but should be %d", ts, chunk.Prev, ccm.chunks[ccm.keys[i-1]].Ts)
				}
			}
		}

		if i == len(ccm.keys)-1 {
			if chunk.Next != 0 {
				return fmt.Errorf("Next of last chunk should be 0, but it's %d", chunk.Next)
			}

			// all checks completed
			break
		}

		var nextChunk *CCacheChunk
		if nextChunk, ok = ccm.chunks[ccm.keys[i+1]]; !ok {
			return fmt.Errorf("Ts %d is in ccm.keys but not in ccm.chunks", ccm.keys[i+1])
		}

		if chunk.Next == 0 {
			if chunk.Ts+chunk.Itgen.Span == nextChunk.Ts {
				return fmt.Errorf("Next of chunk at ts %d is set to 0, but the next chunk is present", ts)
			}
		} else {
			if chunk.Next != nextChunk.Ts {
				return fmt.Errorf("Next of chunk at ts %d is set to %d, but it should be %d", ts, chunk.Next, nextChunk.Ts)
			}
		}
	}
	return nil
}
<commit_msg>verify which chunks should and should not be cached<commit_after>package cache

import (
	"errors"
	"fmt"
	"math/rand"
	"reflect"
	"sort"
	"testing"
	"time"

	"github.com/grafana/metrictank/mdata/cache/accnt"
	"github.com/grafana/metrictank/mdata/chunk"
	"github.com/grafana/metrictank/test"
	"github.com/raintank/schema"
)

func generateChunks(b testing.TB, startAt, count, step uint32) []chunk.IterGen {
	res := make([]chunk.IterGen, 0, count)

	values := make([]uint32, step)
	for t0 := startAt; t0 < startAt+(step*uint32(count)); t0 += step {
		c := getItgen(b, values, t0, true)
		res = append(res, c)
	}
	return res
}

func getCCM() (schema.AMKey, *CCacheMetric) {
	amkey, _ := schema.AMKeyFromString("1.12345678901234567890123456789012")
	ccm := NewCCacheMetric(amkey.MKey)
	return amkey, ccm
}

// TestAddAsc tests adding ascending timestamp chunks individually
func TestAddAsc(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		prev := uint32(1)
		for _, chunk := range chunks {
			ccm.Add(prev, chunk)
			prev = chunk.Ts
		}
	})
}

// TestAddDesc1 tests adding chunks that are all descending
func TestAddDesc1(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		for i := len(chunks) - 1; i >= 0; i-- {
			ccm.Add(0, chunks[i])
		}
	})
}

// TestAddDesc4 tests adding chunks that are globally descending
// but in groups 4 are ascending
func TestAddDesc4(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.Add(0, chunks[2])
		ccm.Add(0, chunks[3])
		ccm.Add(0, chunks[4])
		ccm.Add(0, chunks[5])
		ccm.Add(0, chunks[0])
		ccm.Add(0, chunks[1])
	})
}

// TestAddRange tests adding a contiguous range at once
func TestAddRange(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		prev := uint32(10)
		ccm.AddRange(prev, chunks)
	})
}

// TestAddRangeDesc4 benchmarks adding chunks that are globally descending
// but in groups of 4 are ascending. those groups are added via 1 AddRange.
func TestAddRangeDesc4(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.AddRange(0, chunks[2:6])
		ccm.AddRange(0, chunks[0:2])
	})
}

// test executes the run function
// run should generate chunks and add them to the CCacheMetric however it likes,
// but in a way so that the result will be as expected
func testRun(t *testing.T, run func(*CCacheMetric)) {
	amkey, ccm := getCCM()

	run(ccm)

	res := CCSearchResult{}
	ccm.Search(test.NewContext(), amkey, &res, 25, 45)
	if res.Complete != true {
		t.Fatalf("Expected result to be complete, but it was not")
	}

	if res.Start[0].Ts != 20 {
		t.Fatalf("Expected result to start at 20, but had %d", res.Start[0].Ts)
	}

	if res.Start[len(res.Start)-1].Ts != 40 {
		t.Fatalf("Expected result to start at 40, but had %d", res.Start[len(res.Start)-1].Ts)
	}
}

// BenchmarkAddAsc benchmarks adding ascending timestamp chunks individually
func BenchmarkAddAsc(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	prev := uint32(1)
	b.ResetTimer()
	for _, chunk := range chunks {
		ccm.Add(prev, chunk)
		prev = chunk.Ts
	}
}

// BenchmarkAddDesc1 benchmarks adding chunks that are all descending
func BenchmarkAddDesc1(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 1; i >= 0; i-- {
		ccm.Add(0, chunks[i])
	}
}

// BenchmarkAddDesc4 benchmarks adding chunks that are globally descending
// but in groups 4 are ascending
func BenchmarkAddDesc4(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 4)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 4; i >= 0; i -= 4 {
		ccm.Add(0, chunks[i])
		ccm.Add(0, chunks[i+1])
		ccm.Add(0, chunks[i+2])
		ccm.Add(0, chunks[i+3])
	}
}

// BenchmarkAddDesc64 benchmarks adding chunks that are globally descending
// but in groups 64 are ascending
func BenchmarkAddDesc64(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 64)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 64; i >= 0; i -= 64 {
		for offset := 0; offset < 64; offset += 1 {
			ccm.Add(0, chunks[i+offset])
		}
	}

}

// BenchmarkAddRangeAsc benchmarks adding a contiguous range at once
func BenchmarkAddRangeAsc(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	prev := uint32(1)
	b.ResetTimer()
	ccm.AddRange(prev, chunks)
}

// BenchmarkAddRangeDesc4 benchmarks adding chunks that are globally descending
// but in groups 4 are ascending. those groups are added via 1 AddRange.
func BenchmarkAddRangeDesc4(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 4)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 4; i >= 0; i -= 4 {
		ccm.AddRange(0, chunks[i:i+4])
	}
}

// BenchmarkAddRangeDesc64 benchmarks adding chunks that are globally descending
// but in groups 64 are ascending. those groups are added via 1 AddRange.
func BenchmarkAddRangeDesc64(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 64)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 64; i >= 0; i -= 64 {
		ccm.AddRange(0, chunks[i:i+64])
	}
}

func TestCorruptionCase1(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.AddRange(0, chunks[3:6])
		ccm.AddRange(0, chunks[0:4])
		if err := verifyCcm(ccm, []uint32{10, 20, 30, 40, 50, 60}); err != nil {
			t.Fatal(err)
		}
	})
}

func getRandomNumber(min, max int) int {
	return rand.Intn(max-min) + min
}

// getRandomRange returns a range start-end so that
// end >= start and both numbers drawn from [min, max)
func getRandomRange(min, max int) (int, int) {
	number1 := getRandomNumber(min, max)
	number2 := getRandomNumber(min, max)
	if number1 > number2 {
		return number2, number1
	} else {
		return number1, number2
	}
}

func TestCorruptionCase2(t *testing.T) {
	rand.Seed(time.Now().Unix())
	_, ccm := getCCM()
	iterations := 100000
	var cached [100]bool // tracks which chunks should be cached
	var expKeys []uint32 // tracks which keys should be in the CCM

	// 100 chunks, first t0=10, last is t0=1000
	chunks := generateChunks(t, 10, 100, 10)

	var opAdd, opAddRange, opDel, opDelRange int
	var adds, dels int

	for i := 0; i < iterations; i++ {
		// 0 = Add
		// 1 = AddRange
		// 2 = Del
		// 3 = Del range (via multi del cals)
		action := getRandomNumber(0, 4)
		switch action {
		case 0:
			chunk := getRandomNumber(0, 100)
			t.Logf("adding chunk %d", chunk)
			ccm.Add(0, chunks[chunk])
			cached[chunk] = true
			opAdd++
			adds++
		case 1:
			from, to := getRandomRange(0, 100)
			t.Logf("adding range %d-%d", from, to)
			ccm.AddRange(0, chunks[from:to])
			for chunk := from; chunk < to; chunk++ {
				cached[chunk] = true
			}
			adds += (to - from)
			opAddRange++
		case 2:
			chunk := getRandomNumber(0, 100)
			t.Logf("deleting chunk %d", chunk)
			ccm.Del(chunks[chunk].Ts) // note: chunk may not exist
			cached[chunk] = false
			opDel++
			dels++
		case 3:
			from, to := getRandomRange(0, 100)
			t.Logf("deleting range %d-%d", from, to)
			for chunk := from; chunk < to; chunk++ {
				ccm.Del(chunks[chunk].Ts) // note: chunk may not exist
				cached[chunk] = false
			}
			opDelRange++
			dels += (to - from)
		}

		expKeys = expKeys[:0]
		for i, c := range cached {
			if c {
				expKeys = append(expKeys, uint32((i+1)*10))
			}
		}

		if err := verifyCcm(ccm, expKeys); err != nil {
			t.Fatal(err)
		}
	}

	fmt.Printf("operations: add %d - addRange %d - del %d - delRange %d\n", opAdd, opAddRange, opDel, opDelRange)
	fmt.Printf("total chunk adds %d - total chunk deletes %d\n", adds, dels)
}

// verifyCcm verifies the integrity of a CCacheMetric
// it assumes that all itergens are span-aware
func verifyCcm(ccm *CCacheMetric, expKeys []uint32) error {
	var chunk *CCacheChunk
	var ok bool

	if len(ccm.chunks) != len(ccm.keys) {
		return errors.New("Length of ccm.chunks does not match ccm.keys")
	}

	if !sort.IsSorted(accnt.Uint32Asc(ccm.keys)) {
		return errors.New("keys are not sorted")
	}

	if !reflect.DeepEqual(ccm.keys, expKeys) {
		return fmt.Errorf("keys mismatch. expected %v, got %v", expKeys, ccm.keys)
	}

	for i, ts := range ccm.keys {
		if chunk, ok = ccm.chunks[ts]; !ok {
			return fmt.Errorf("Ts %d is in ccm.keys but not in ccm.chunks", ts)
		}

		if i == 0 {
			if chunk.Prev != 0 {
				return errors.New("First chunk has Prev != 0")
			}
		} else {
			if chunk.Prev == 0 {
				if ccm.chunks[ccm.keys[i-1]].Ts == chunk.Ts-chunk.Itgen.Span {
					return fmt.Errorf("Chunk of ts %d has Prev == 0, but the previous chunk is present", ts)
				}
			} else {
				if ccm.chunks[ccm.keys[i-1]].Ts != chunk.Prev {
					return fmt.Errorf("Chunk of ts %d has Prev set to wrong ts %d but should be %d", ts, chunk.Prev, ccm.chunks[ccm.keys[i-1]].Ts)
				}
			}
		}

		if i == len(ccm.keys)-1 {
			if chunk.Next != 0 {
				return fmt.Errorf("Next of last chunk should be 0, but it's %d", chunk.Next)
			}

			// all checks completed
			break
		}

		var nextChunk *CCacheChunk
		if nextChunk, ok = ccm.chunks[ccm.keys[i+1]]; !ok {
			return fmt.Errorf("Ts %d is in ccm.keys but not in ccm.chunks", ccm.keys[i+1])
		}

		if chunk.Next == 0 {
			if chunk.Ts+chunk.Itgen.Span == nextChunk.Ts {
				return fmt.Errorf("Next of chunk at ts %d is set to 0, but the next chunk is present", ts)
			}
		} else {
			if chunk.Next != nextChunk.Ts {
				return fmt.Errorf("Next of chunk at ts %d is set to %d, but it should be %d", ts, chunk.Next, nextChunk.Ts)
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package metrics

import (
	"regexp"

	"github.com/czerwonk/bird_exporter/protocol"
)

// DefaultLabelStrategy defines the labels to add to an metric and its data retrieval method
type DefaultLabelStrategy struct {
	descriptionLabels      bool
	descriptionLabelsRegex string
}

func NewDefaultLabelStrategy(descriptionLabels bool, descriptionLabelsRegex string) *DefaultLabelStrategy {
	return &DefaultLabelStrategy{
		descriptionLabels:      descriptionLabels,
		descriptionLabelsRegex: descriptionLabelsRegex,
	}
}

// LabelNames returns the list of label names
func (d *DefaultLabelStrategy) LabelNames(p *protocol.Protocol) []string {
	res := []string{"name", "proto", "ip_version", "import_filter", "export_filter"}
	if d.descriptionLabels && p.Description != "" {
		res = append(res, labelKeysFromDescription(p.Description, d)...)
	}

	return res
}

// LabelValues returns the values for a protocol
func (d *DefaultLabelStrategy) LabelValues(p *protocol.Protocol) []string {
	res := []string{p.Name, protoString(p), p.IPVersion, p.ImportFilter, p.ExportFilter}
	if d.descriptionLabels && p.Description != "" {
		res = append(res, labelValuesFromDescription(p.Description, d)...)
	}

	return res
}

func labelKeysFromDescription(desc string, d *DefaultLabelStrategy) (res []string) {
	reAllStringSubmatch := labelFindAllStringSubmatch(desc, d)
	for _, submatch := range reAllStringSubmatch {
		res = append(res, submatch[1])
	}

	return
}

func labelValuesFromDescription(desc string, d *DefaultLabelStrategy) (res []string) {
	reAllStringSubmatch := labelFindAllStringSubmatch(desc, d)
	for _, submatch := range reAllStringSubmatch {
		res = append(res, submatch[2])
	}

	return
}

func labelFindAllStringSubmatch(desc string, d *DefaultLabelStrategy) (result [][]string) {

	// Regex pattern captures "key: value" pair from the content.
	pattern := regexp.MustCompile(d.descriptionLabelsRegex)

	result = pattern.FindAllStringSubmatch(desc, -1)

	return
}

func protoString(p *protocol.Protocol) string {
	switch p.Proto {
	case protocol.BGP:
		return "BGP"
	case protocol.OSPF:
		if p.IPVersion == "4" {
			return "OSPF"
		}
		return "OSPFv3"
	case protocol.Static:
		return "Static"
	case protocol.Kernel:
		return "Kernel"
	case protocol.Direct:
		return "Direct"
	case protocol.Babel:
		return "Babel"
	case protocol.RPKI:
		return "RPKI"
	case protocol.BFD:
		return "BFD"
	}

	return ""
}
<commit_msg>reuse regex, ensure custom label names are trimmed<commit_after>package metrics

import (
	"regexp"
	"strings"

	"github.com/czerwonk/bird_exporter/protocol"
)

// DefaultLabelStrategy defines the labels to add to an metric and its data retrieval method
type DefaultLabelStrategy struct {
	descriptionLabels      bool
	descriptionLabelsRegex *regexp.Regexp
}

func NewDefaultLabelStrategy(descriptionLabels bool, descriptionLabelsRegex string) *DefaultLabelStrategy {
	return &DefaultLabelStrategy{
		descriptionLabels:      descriptionLabels,
		descriptionLabelsRegex: regexp.MustCompile(descriptionLabelsRegex),
	}
}

// LabelNames returns the list of label names
func (d *DefaultLabelStrategy) LabelNames(p *protocol.Protocol) []string {
	res := []string{"name", "proto", "ip_version", "import_filter", "export_filter"}
	if d.descriptionLabels && p.Description != "" {
		res = append(res, labelKeysFromDescription(p.Description, d)...)
	}

	return res
}

// LabelValues returns the values for a protocol
func (d *DefaultLabelStrategy) LabelValues(p *protocol.Protocol) []string {
	res := []string{p.Name, protoString(p), p.IPVersion, p.ImportFilter, p.ExportFilter}
	if d.descriptionLabels && p.Description != "" {
		res = append(res, labelValuesFromDescription(p.Description, d)...)
	}

	return res
}

func labelKeysFromDescription(desc string, d *DefaultLabelStrategy) []string {
	res := []string{}

	matches := d.descriptionLabelsRegex.FindAllStringSubmatch(desc, -1)
	for _, submatch := range matches {
		res = append(res, strings.TrimSpace(submatch[1]))
	}

	return res
}

func labelValuesFromDescription(desc string, d *DefaultLabelStrategy) []string {
	res := []string{}

	matches := d.descriptionLabelsRegex.FindAllStringSubmatch(desc, -1)
	for _, submatch := range matches {
		res = append(res, strings.TrimSpace(submatch[2]))
	}

	return res
}

func protoString(p *protocol.Protocol) string {
	switch p.Proto {
	case protocol.BGP:
		return "BGP"
	case protocol.OSPF:
		if p.IPVersion == "4" {
			return "OSPF"
		}
		return "OSPFv3"
	case protocol.Static:
		return "Static"
	case protocol.Kernel:
		return "Kernel"
	case protocol.Direct:
		return "Direct"
	case protocol.Babel:
		return "Babel"
	case protocol.RPKI:
		return "RPKI"
	case protocol.BFD:
		return "BFD"
	}

	return ""
}
<|endoftext|>


<commit_before>package xsrf

import (
	"fmt"
	"html/template"
	"net/http"
	"time"

	"github.com/go-xweb/uuid"
	"github.com/lunny/tango"
)

const (
	XSRF_TAG string = "_xsrf"
)

func XsrfName() string {
	return XSRF_TAG
}

type Xsrfer interface {
	CheckXsrf() bool
}

type NoCheck struct {
}

func (NoCheck) CheckXsrf() bool {
	return false
}

type XsrfChecker interface {
	SetXsrfValue(string)
	SetXsrfFormHtml(template.HTML)
}

type Checker struct {
	XsrfValue string
	XsrfFormHtml template.HTML
}

func (Checker) CheckXsrf() bool {
	return true
}

func (c *Checker) SetXsrfValue(v string) {
	c.XsrfValue = v
}

func (c *Checker) SetXsrfFormHtml(t template.HTML) {
	c.XsrfFormHtml = t
}

type Xsrf struct {
	timeout time.Duration
}

func New(timeout time.Duration) *Xsrf {
	return &Xsrf{
		timeout: timeout,
	}
}

// NewCookie is a helper method that returns a new http.Cookie object.
// Duration is specified in seconds. If the duration is zero, the cookie is permanent.
// This can be used in conjunction with ctx.SetCookie.
func newCookie(name string, value string, age int64) *http.Cookie {
	var utctime time.Time
	if age == 0 {
		// 2^31 - 1 seconds (roughly 2038)
		utctime = time.Unix(2147483647, 0)
	} else {
		utctime = time.Unix(time.Now().Unix()+age, 0)
	}
	return &http.Cookie{Name: name, Value: value, Expires: utctime}
}

func (xsrf *Xsrf) Handle(ctx *tango.Context) {
	var action interface{}
	if action = ctx.Action(); action == nil {
		ctx.Next()
		return
	}

	// if action implements check xsrf option and ask not check then return
	if checker, ok := action.(Xsrfer); ok && !checker.CheckXsrf() {
		ctx.Next()
		return
	}

	if ctx.Req().Method == "GET" {
		var val string = ""
		cookie, err := ctx.Req().Cookie(XSRF_TAG)
		if err != nil {
			val = uuid.NewRandom().String()
			cookie = newCookie(XSRF_TAG, val, int64(xsrf.timeout))
			ctx.Header().Set("Set-Cookie", cookie.String())
		} else {
			val = cookie.Value
		}

		if c, ok := action.(XsrfChecker); ok {
			c.SetXsrfValue(val)
			c.SetXsrfFormHtml(template.HTML(fmt.Sprintf(`<input type="hidden" name="%v" value="%v" />`,
				XSRF_TAG, val)))
		}
	} else if ctx.Req().Method == "POST" {
		res, err := ctx.Req().Cookie(XSRF_TAG)
		formVals := ctx.Req().Form[XSRF_TAG]
		var formVal string
		if len(formVals) > 0 {
			formVal = formVals[0]
		}
		if err != nil || res.Value == "" || res.Value != formVal {
			ctx.WriteHeader(http.StatusInternalServerError)
			ctx.Write([]byte("xsrf token error."))
			return
		}
	}

	ctx.Next()
}
<commit_msg>add log<commit_after>package xsrf

import (
	"fmt"
	"html/template"
	"net/http"
	"time"

	"github.com/go-xweb/uuid"
	"github.com/lunny/tango"
)

const (
	XSRF_TAG string = "_xsrf"
)

func XsrfName() string {
	return XSRF_TAG
}

type Xsrfer interface {
	CheckXsrf() bool
}

type NoCheck struct {
}

func (NoCheck) CheckXsrf() bool {
	return false
}

type XsrfChecker interface {
	SetXsrfValue(string)
	SetXsrfFormHtml(template.HTML)
}

type Checker struct {
	XsrfValue string
	XsrfFormHtml template.HTML
}

func (Checker) CheckXsrf() bool {
	return true
}

func (c *Checker) SetXsrfValue(v string) {
	c.XsrfValue = v
}

func (c *Checker) SetXsrfFormHtml(t template.HTML) {
	c.XsrfFormHtml = t
}

type Xsrf struct {
	timeout time.Duration
}

func New(timeout time.Duration) *Xsrf {
	return &Xsrf{
		timeout: timeout,
	}
}

// NewCookie is a helper method that returns a new http.Cookie object.
// Duration is specified in seconds. If the duration is zero, the cookie is permanent.
// This can be used in conjunction with ctx.SetCookie.
func newCookie(name string, value string, age int64) *http.Cookie {
	var utctime time.Time
	if age == 0 {
		// 2^31 - 1 seconds (roughly 2038)
		utctime = time.Unix(2147483647, 0)
	} else {
		utctime = time.Unix(time.Now().Unix()+age, 0)
	}
	return &http.Cookie{Name: name, Value: value, Expires: utctime}
}

func (xsrf *Xsrf) Handle(ctx *tango.Context) {
	var action interface{}
	if action = ctx.Action(); action == nil {
		ctx.Next()
		return
	}

	// if action implements check xsrf option and ask not check then return
	if checker, ok := action.(Xsrfer); ok && !checker.CheckXsrf() {
		ctx.Next()
		return
	}

	if ctx.Req().Method == "GET" {
		var val string = ""
		cookie, err := ctx.Req().Cookie(XSRF_TAG)
		if err != nil {
			val = uuid.NewRandom().String()
			cookie = newCookie(XSRF_TAG, val, int64(xsrf.timeout))
			ctx.Header().Set("Set-Cookie", cookie.String())
		} else {
			val = cookie.Value
		}

		if c, ok := action.(XsrfChecker); ok {
			c.SetXsrfValue(val)
			c.SetXsrfFormHtml(template.HTML(fmt.Sprintf(`<input type="hidden" name="%v" value="%v" />`,
				XSRF_TAG, val)))
		}

	} else if ctx.Req().Method == "POST" {
		res, err := ctx.Req().Cookie(XSRF_TAG)
		formVal := ctx.Req().FormValue(XSRF_TAG)

		if err != nil || res.Value == "" || res.Value != formVal {
			ctx.Abort(http.StatusInternalServerError, "xsrf token error.")
			ctx.Error("xsrf token error.")
			return
		}
	}

	ctx.Next()
}
<|endoftext|>


<commit_before>package sarama

import (
	"sync"

	"github.com/klauspost/compress/zstd"
)

var (
	zstdDec *zstd.Decoder
	zstdEnc *zstd.Encoder

	zstdEncOnce, zstdDecOnce sync.Once
)

func zstdDecompress(dst, src []byte) ([]byte, error) {
	zstdDecOnce.Do(func() {
		zstdDec, _ = zstd.NewReader(nil)
	})
	return zstdDec.DecodeAll(src, dst)
}

func zstdCompress(dst, src []byte) ([]byte, error) {
	zstdEncOnce.Do(func() {
		zstdEnc, _ = zstd.NewWriter(nil, zstd.WithZeroFrames(true))
	})
	return zstdEnc.EncodeAll(src, dst), nil
}
<commit_msg>feat: zstd performance improvements<commit_after>package sarama

import (
	"github.com/klauspost/compress/zstd"
)

var zstdDec, _ = zstd.NewReader(nil)
var zstdEnc, _ = zstd.NewWriter(nil, zstd.WithZeroFrames(true))

func zstdDecompress(dst, src []byte) ([]byte, error) {
	return zstdDec.DecodeAll(src, dst)
}

func zstdCompress(dst, src []byte) ([]byte, error) {
	return zstdEnc.EncodeAll(src, dst), nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. 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 gcs

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"math/rand"
	"net"
	"net/url"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/api/googleapi"
)

// A bucket that wraps another, calling its methods in a retry loop with
// randomized exponential backoff.
type retryBucket struct {
	maxSleep time.Duration
	wrapped  Bucket
}

func newRetryBucket(
		maxSleep time.Duration,
		wrapped Bucket) (b Bucket) {
	b = &retryBucket{
		maxSleep: maxSleep,
		wrapped:  wrapped,
	}

	return
}

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

func shouldRetry(err error) (b bool) {
	// HTTP 50x errors.
	if typed, ok := err.(*googleapi.Error); ok {
		if typed.Code >= 500 && typed.Code < 600 {
			b = true
			return
		}
	}

	// HTTP 429 errors (GCS uses these for rate limiting).
	if typed, ok := err.(*googleapi.Error); ok {
		if typed.Code == 429 {
			b = true
			return
		}
	}

	// Network errors, which tend to show up transiently when doing lots of
	// operations in parallel. For example:
	//
	//     dial tcp 74.125.203.95:443: too many open files
	//
	if _, ok := err.(*net.OpError); ok {
		b = true
		return
	}

	// The HTTP package returns ErrUnexpectedEOF in several places. This seems to
	// come up when the server terminates the connection in the middle of an
	// object read.
	if err == io.ErrUnexpectedEOF {
		b = true
		return
	}

	// The HTTP library also appears to leak EOF errors from... somewhere in its
	// guts as URL errors sometimes.
	if urlErr, ok := err.(*url.Error); ok {
		if urlErr.Err == io.EOF {
			b = true
			return
		}
	}

	// Sometimes the HTTP package helpfully encapsulates the real error in a URL
	// error.
	if urlErr, ok := err.(*url.Error); ok {
		b = shouldRetry(urlErr.Err)
		return
	}

	return
}

// Choose an appropriate delay for exponential backoff, given that we have
// already slept the given number of times for this logical request.
func chooseDelay(prevSleepCount uint) (d time.Duration) {
	const baseDelay = time.Millisecond

	// Choose a a delay in [0, 2^prevSleepCount * baseDelay).
	d = (1 << prevSleepCount) * baseDelay
	d = time.Duration(rand.Int63n(int64(d)))

	return
}

// Exponential backoff for a function that might fail.
//
// This is essentially what is described in the "Best practices" section of the
// "Upload Objects" docs:
//
//	https://cloud.google.com/storage/docs/json_api/v1/how-tos/upload
//
// with the following exceptions:
//
//   - We perform backoff for all operations.
//
//   - The random component scales with the delay, so that the first sleep
//     cannot be as long as one second. The algorithm used matches the
//     description at http://en.wikipedia.org/wiki/Exponential_backoff.
//
//   - We retry more types of errors; see shouldRetry above.
//
// State for total sleep time and number of previous sleeps is housed outside
// of this function to allow it to be "resumed" by multiple invocations of
// retryObjectReader.Read.
func expBackoff(
		ctx context.Context,
		desc string,
		maxSleep time.Duration,
		f func() error,
		prevSleepCount *uint,
		prevSleepDuration *time.Duration) (err error) {
	for {
		// Make an attempt. Stop if successful.
		err = f()
		if err == nil {
			return
		}

		// Do we want to retry?
		if !shouldRetry(err) {
			// Special case: don't spam up the logs for EOF, which io.Reader returns
			// in the normal course of things.
			if err != io.EOF {
				err = fmt.Errorf("not retrying %s: %w", desc, err)
			}

			return
		}

		// Choose a a delay.
		d := chooseDelay(*prevSleepCount)
		*prevSleepCount++

		// Are we out of credit?
		if *prevSleepDuration+d > maxSleep {
			// Return the most recent error.
			return
		}

		// Sleep, returning early if cancelled.
		select {
		case <-ctx.Done():
			// On cancellation, return the last error we saw.
			return

		case <-time.After(d):
			*prevSleepDuration += d
			continue
		}
	}
}

// Like expBackoff, but assumes that we've never slept before (and won't need
// to sleep again).
func oneShotExpBackoff(
		ctx context.Context,
		desc string,
		maxSleep time.Duration,
		f func() error) (err error) {
	var prevSleepCount uint
	var prevSleepDuration time.Duration

	err = expBackoff(
		ctx,
		desc,
		maxSleep,
		f,
		&prevSleepCount,
		&prevSleepDuration)

	return
}

////////////////////////////////////////////////////////////////////////
// Read support
////////////////////////////////////////////////////////////////////////

type retryObjectReader struct {
	bucket *retryBucket

	// The context we should watch when sleeping for retries.
	ctx context.Context

	// What we are trying to read.
	name       string
	generation int64
	byteRange  ByteRange

	// nil when we start or have seen a permanent error.
	wrapped io.ReadCloser

	// If we've seen an error that we shouldn't retry for, this will be non-nil
	// and should be returned permanently.
	permanentErr error

	// The number of times we've slept so far, and the total amount of time we've
	// spent sleeping.
	sleepCount    uint
	sleepDuration time.Duration
}

// Set up the wrapped reader.
func (rc *retryObjectReader) setUpWrapped() (err error) {
	// Call through to create the reader.
	req := &ReadObjectRequest{
		Name:       rc.name,
		Generation: rc.generation,
		Range:      &rc.byteRange,
	}

	wrapped, err := rc.bucket.wrapped.NewReader(rc.ctx, req)
	if err != nil {
		return
	}

	rc.wrapped = wrapped
	return
}

// Set up the wrapped reader if necessary, and make one attempt to read through
// it.
//
// Clears the wrapped reader on error.
func (rc *retryObjectReader) readOnce(p []byte) (n int, err error) {
	// Set up the wrapped reader if it's not already around.
	if rc.wrapped == nil {
		err = rc.setUpWrapped()
		if err != nil {
			return
		}
	}

	// Attempt to read from it.
	n, err = rc.wrapped.Read(p)
	if err != nil {
		rc.wrapped.Close()
		rc.wrapped = nil
		return
	}

	return
}

// Invariant: we never return an error from this function unless we've given up
// on retrying. In particular, we won't return a short read because the wrapped
// reader returned a short read and an error.
func (rc *retryObjectReader) Read(p []byte) (n int, err error) {
	// Whatever we do, accumulate the bytes that we're returning to the user.
	defer func() {
		if n < 0 {
			panic(fmt.Sprintf("Negative byte count: %d", n))
		}

		rc.byteRange.Start += uint64(n)
	}()

	// If we've already decided on a permanent error, return that.
	if rc.permanentErr != nil {
		err = rc.permanentErr
		return
	}

	// If we let an error escape below, it must be a permanent one.
	defer func() {
		if err != nil {
			rc.permanentErr = err
		}
	}()

	// We will repeatedly make single attempts until we get a successful request.
	// Don't forget to accumulate the result each time.
	tryOnce := func() (err error) {
		var bytesRead int
		bytesRead, err = rc.readOnce(p)
		n += bytesRead
		p = p[bytesRead:]

		return
	}

	err = expBackoff(
		rc.ctx,
		fmt.Sprintf("Read(%q, %d)", rc.name, rc.generation),
		rc.bucket.maxSleep,
		tryOnce,
		&rc.sleepCount,
		&rc.sleepDuration)

	return
}

func (rc *retryObjectReader) Close() (err error) {
	// If we don't have a wrapped reader, there is nothing useful that we can or
	// need to do here.
	if rc.wrapped == nil {
		return
	}

	// Call through.
	err = rc.wrapped.Close()

	return
}

////////////////////////////////////////////////////////////////////////
// Public interface
////////////////////////////////////////////////////////////////////////

func (rb *retryBucket) Name() (name string) {
	name = rb.wrapped.Name()
	return
}

func (rb *retryBucket) NewReader(
		ctx context.Context,
		req *ReadObjectRequest) (rc io.ReadCloser, err error) {
	// If the user specified the latest generation, we need to figure out what
	// that is so that we can create a reader that knows how to keep a stable
	// generation despite retrying repeatedly.
	var generation int64 = req.Generation
	var sleepCount uint
	var sleepDuration time.Duration

	if generation == 0 {
		findGeneration := func() (err error) {
			o, err := rb.wrapped.StatObject(
				ctx,
				&StatObjectRequest{
					Name: req.Name,
				})

			if err != nil {
				return
			}

			generation = o.Generation
			return
		}

		err = expBackoff(
			ctx,
			fmt.Sprintf("FindLatestGeneration(%q)", req.Name),
			rb.maxSleep,
			findGeneration,
			&sleepCount,
			&sleepDuration)

		if err != nil {
			return
		}
	}

	// Choose an appropriate byte range.
	byteRange := ByteRange{0, math.MaxUint64}
	if req.Range != nil {
		byteRange = *req.Range
	}

	// Now that we know what generation we're looking for, return an appropriate
	// reader that knows how to retry when the connection fails. Make sure to
	// inherit the time spent sleeping above.
	rc = &retryObjectReader{
		bucket: rb,
		ctx:    ctx,

		name:       req.Name,
		generation: generation,
		byteRange:  byteRange,

		sleepCount:    sleepCount,
		sleepDuration: sleepDuration,
	}

	return
}

func (rb *retryBucket) CreateObject(
		ctx context.Context,
		req *CreateObjectRequest) (o *Object, err error) {
	var seeker io.ReadSeeker
	if readSeeker, ok := req.Contents.(io.ReadSeeker); ok {
		seeker = readSeeker
	} else {
		// We can't simply replay the request multiple times, because the first
		// attempt might exhaust some of the req.Contents reader, leaving
		// missing contents for the second attempt.
		//
		// So, copy out all contents and create a copy of the request that we
		// will modify to serve from memory for each call.
		data, err := ioutil.ReadAll(req.Contents)
		if err != nil {
			err = fmt.Errorf("ioutil.ReadAll: %v", err)
			return nil, err
		}
		seeker = bytes.NewReader(data)
	}

	reqCopy := *req

	// Call through with that request.
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("CreateObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			if _, err = seeker.Seek(0, io.SeekStart); err != nil {
				return
			}
			reqCopy.Contents = seeker
			o, err = rb.wrapped.CreateObject(ctx, &reqCopy)
			return
		})

	return
}

func (rb *retryBucket) CopyObject(
		ctx context.Context,
		req *CopyObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("CopyObject(%q, %q)", req.SrcName, req.DstName),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.CopyObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) ComposeObjects(
		ctx context.Context,
		req *ComposeObjectsRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("ComposeObjects(%q)", req.DstName),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.ComposeObjects(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) StatObject(
		ctx context.Context,
		req *StatObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("StatObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.StatObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) ListObjects(
		ctx context.Context,
		req *ListObjectsRequest) (listing *Listing, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("ListObjects(%q)", req.Prefix),
		rb.maxSleep,
		func() (err error) {
			listing, err = rb.wrapped.ListObjects(ctx, req)
			return
		})
	return
}

func (rb *retryBucket) UpdateObject(
		ctx context.Context,
		req *UpdateObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("UpdateObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.UpdateObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) DeleteObject(
		ctx context.Context,
		req *DeleteObjectRequest) (err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("DeleteObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			err = rb.wrapped.DeleteObject(ctx, req)
			return
		})

	return
}
<commit_msg>Unit Test for BucketMnager<commit_after>// Copyright 2015 Google Inc. 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 gcs

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"math/rand"
	"net"
	"net/url"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/api/googleapi"
)

// A bucket that wraps another, calling its methods in a retry loop with
// randomized exponential backoff.
type retryBucket struct {
	maxSleep time.Duration
	wrapped  Bucket
}

func newRetryBucket(
		maxSleep time.Duration,
		wrapped Bucket) (b Bucket) {
	b = &retryBucket{
		maxSleep: maxSleep,
		wrapped:  wrapped,
	}

	return
}

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

func shouldRetry(err error) (b bool) {
	// HTTP 50x errors.
	if typed, ok := err.(*googleapi.Error); ok {
		if typed.Code >= 500 && typed.Code < 600 {
			b = true
			return
		}
	}

	// HTTP 429 errors (GCS uses these for rate limiting).
	if typed, ok := err.(*googleapi.Error); ok {
		if typed.Code == 429 {
			b = true
			return
		}
	}

	// Network errors, which tend to show up transiently when doing lots of
	// operations in parallel. For example:
	//
	//     dial tcp 74.125.203.95:443: too many open files
	//
	if _, ok := err.(*net.OpError); ok {
		b = true
		return
	}

	// The HTTP package returns ErrUnexpectedEOF in several places. This seems to
	// come up when the server terminates the connection in the middle of an
	// object read.
	if err == io.ErrUnexpectedEOF {
		b = true
		return
	}

	// The HTTP library also appears to leak EOF errors from... somewhere in its
	// guts as URL errors sometimes.
	if urlErr, ok := err.(*url.Error); ok {
		if urlErr.Err == io.EOF {
			b = true
			return
		}
	}

	// Sometimes the HTTP package helpfully encapsulates the real error in a URL
	// error.
	if urlErr, ok := err.(*url.Error); ok {
		b = shouldRetry(urlErr.Err)
		return
	}

	return
}

// Choose an appropriate delay for exponential backoff, given that we have
// already slept the given number of times for this logical request.
func chooseDelay(prevSleepCount uint) (d time.Duration) {
	const baseDelay = time.Millisecond

	// Choose a a delay in [0, 2^prevSleepCount * baseDelay).
	d = (1 << prevSleepCount) * baseDelay
	d = time.Duration(rand.Int63n(int64(d)))

	return
}

// Exponential backoff for a function that might fail.
//
// This is essentially what is described in the "Best practices" section of the
// "Upload Objects" docs:
//
//     https://cloud.google.com/storage/docs/json_api/v1/how-tos/upload
//
// with the following exceptions:
//
//  *  We perform backoff for all operations.
//
//  *  The random component scales with the delay, so that the first sleep
//     cannot be as long as one second. The algorithm used matches the
//     description at http://en.wikipedia.org/wiki/Exponential_backoff.
//
//  *  We retry more types of errors; see shouldRetry above.
//
// State for total sleep time and number of previous sleeps is housed outside
// of this function to allow it to be "resumed" by multiple invocations of
// retryObjectReader.Read.
func expBackoff(
		ctx context.Context,
		desc string,
		maxSleep time.Duration,
		f func() error,
		prevSleepCount *uint,
		prevSleepDuration *time.Duration) (err error) {
	for {
		// Make an attempt. Stop if successful.
		err = f()
		if err == nil {
			return
		}

		// Do we want to retry?
		if !shouldRetry(err) {
			// Special case: don't spam up the logs for EOF, which io.Reader returns
			// in the normal course of things.
			if err != io.EOF {
				err = fmt.Errorf("not retrying %s: %w", desc, err)
			}

			return
		}

		// Choose a a delay.
		d := chooseDelay(*prevSleepCount)
		*prevSleepCount++

		// Are we out of credit?
		if *prevSleepDuration+d > maxSleep {
			// Return the most recent error.
			return
		}

		// Sleep, returning early if cancelled.
		select {
		case <-ctx.Done():
			// On cancellation, return the last error we saw.
			return

		case <-time.After(d):
			*prevSleepDuration += d
			continue
		}
	}
}

// Like expBackoff, but assumes that we've never slept before (and won't need
// to sleep again).
func oneShotExpBackoff(
		ctx context.Context,
		desc string,
		maxSleep time.Duration,
		f func() error) (err error) {
	var prevSleepCount uint
	var prevSleepDuration time.Duration

	err = expBackoff(
		ctx,
		desc,
		maxSleep,
		f,
		&prevSleepCount,
		&prevSleepDuration)

	return
}

////////////////////////////////////////////////////////////////////////
// Read support
////////////////////////////////////////////////////////////////////////

type retryObjectReader struct {
	bucket *retryBucket

	// The context we should watch when sleeping for retries.
	ctx context.Context

	// What we are trying to read.
	name       string
	generation int64
	byteRange  ByteRange

	// nil when we start or have seen a permanent error.
	wrapped io.ReadCloser

	// If we've seen an error that we shouldn't retry for, this will be non-nil
	// and should be returned permanently.
	permanentErr error

	// The number of times we've slept so far, and the total amount of time we've
	// spent sleeping.
	sleepCount    uint
	sleepDuration time.Duration
}

// Set up the wrapped reader.
func (rc *retryObjectReader) setUpWrapped() (err error) {
	// Call through to create the reader.
	req := &ReadObjectRequest{
		Name:       rc.name,
		Generation: rc.generation,
		Range:      &rc.byteRange,
	}

	wrapped, err := rc.bucket.wrapped.NewReader(rc.ctx, req)
	if err != nil {
		return
	}

	rc.wrapped = wrapped
	return
}

// Set up the wrapped reader if necessary, and make one attempt to read through
// it.
//
// Clears the wrapped reader on error.
func (rc *retryObjectReader) readOnce(p []byte) (n int, err error) {
	// Set up the wrapped reader if it's not already around.
	if rc.wrapped == nil {
		err = rc.setUpWrapped()
		if err != nil {
			return
		}
	}

	// Attempt to read from it.
	n, err = rc.wrapped.Read(p)
	if err != nil {
		rc.wrapped.Close()
		rc.wrapped = nil
		return
	}

	return
}

// Invariant: we never return an error from this function unless we've given up
// on retrying. In particular, we won't return a short read because the wrapped
// reader returned a short read and an error.
func (rc *retryObjectReader) Read(p []byte) (n int, err error) {
	// Whatever we do, accumulate the bytes that we're returning to the user.
	defer func() {
		if n < 0 {
			panic(fmt.Sprintf("Negative byte count: %d", n))
		}

		rc.byteRange.Start += uint64(n)
	}()

	// If we've already decided on a permanent error, return that.
	if rc.permanentErr != nil {
		err = rc.permanentErr
		return
	}

	// If we let an error escape below, it must be a permanent one.
	defer func() {
		if err != nil {
			rc.permanentErr = err
		}
	}()

	// We will repeatedly make single attempts until we get a successful request.
	// Don't forget to accumulate the result each time.
	tryOnce := func() (err error) {
		var bytesRead int
		bytesRead, err = rc.readOnce(p)
		n += bytesRead
		p = p[bytesRead:]

		return
	}

	err = expBackoff(
		rc.ctx,
		fmt.Sprintf("Read(%q, %d)", rc.name, rc.generation),
		rc.bucket.maxSleep,
		tryOnce,
		&rc.sleepCount,
		&rc.sleepDuration)

	return
}

func (rc *retryObjectReader) Close() (err error) {
	// If we don't have a wrapped reader, there is nothing useful that we can or
	// need to do here.
	if rc.wrapped == nil {
		return
	}

	// Call through.
	err = rc.wrapped.Close()

	return
}

////////////////////////////////////////////////////////////////////////
// Public interface
////////////////////////////////////////////////////////////////////////

func (rb *retryBucket) Name() (name string) {
	name = rb.wrapped.Name()
	return
}

func (rb *retryBucket) NewReader(
		ctx context.Context,
		req *ReadObjectRequest) (rc io.ReadCloser, err error) {
	// If the user specified the latest generation, we need to figure out what
	// that is so that we can create a reader that knows how to keep a stable
	// generation despite retrying repeatedly.
	var generation int64 = req.Generation
	var sleepCount uint
	var sleepDuration time.Duration

	if generation == 0 {
		findGeneration := func() (err error) {
			o, err := rb.wrapped.StatObject(
				ctx,
				&StatObjectRequest{
					Name: req.Name,
				})

			if err != nil {
				return
			}

			generation = o.Generation
			return
		}

		err = expBackoff(
			ctx,
			fmt.Sprintf("FindLatestGeneration(%q)", req.Name),
			rb.maxSleep,
			findGeneration,
			&sleepCount,
			&sleepDuration)

		if err != nil {
			return
		}
	}

	// Choose an appropriate byte range.
	byteRange := ByteRange{0, math.MaxUint64}
	if req.Range != nil {
		byteRange = *req.Range
	}

	// Now that we know what generation we're looking for, return an appropriate
	// reader that knows how to retry when the connection fails. Make sure to
	// inherit the time spent sleeping above.
	rc = &retryObjectReader{
		bucket: rb,
		ctx:    ctx,

		name:       req.Name,
		generation: generation,
		byteRange:  byteRange,

		sleepCount:    sleepCount,
		sleepDuration: sleepDuration,
	}

	return
}

func (rb *retryBucket) CreateObject(
		ctx context.Context,
		req *CreateObjectRequest) (o *Object, err error) {
	var seeker io.ReadSeeker
	if readSeeker, ok := req.Contents.(io.ReadSeeker); ok {
		seeker = readSeeker
	} else {
		// We can't simply replay the request multiple times, because the first
		// attempt might exhaust some of the req.Contents reader, leaving
		// missing contents for the second attempt.
		//
		// So, copy out all contents and create a copy of the request that we
		// will modify to serve from memory for each call.
		data, err := ioutil.ReadAll(req.Contents)
		if err != nil {
			err = fmt.Errorf("ioutil.ReadAll: %v", err)
			return nil, err
		}
		seeker = bytes.NewReader(data)
	}

	reqCopy := *req

	// Call through with that request.
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("CreateObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			if _, err = seeker.Seek(0, io.SeekStart); err != nil {
				return
			}
			reqCopy.Contents = seeker
			o, err = rb.wrapped.CreateObject(ctx, &reqCopy)
			return
		})

	return
}

func (rb *retryBucket) CopyObject(
		ctx context.Context,
		req *CopyObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("CopyObject(%q, %q)", req.SrcName, req.DstName),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.CopyObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) ComposeObjects(
		ctx context.Context,
		req *ComposeObjectsRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("ComposeObjects(%q)", req.DstName),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.ComposeObjects(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) StatObject(
		ctx context.Context,
		req *StatObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("StatObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.StatObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) ListObjects(
		ctx context.Context,
		req *ListObjectsRequest) (listing *Listing, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("ListObjects(%q)", req.Prefix),
		rb.maxSleep,
		func() (err error) {
			listing, err = rb.wrapped.ListObjects(ctx, req)
			return
		})
	return
}

func (rb *retryBucket) UpdateObject(
		ctx context.Context,
		req *UpdateObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("UpdateObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.UpdateObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) DeleteObject(
		ctx context.Context,
		req *DeleteObjectRequest) (err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("DeleteObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			err = rb.wrapped.DeleteObject(ctx, req)
			return
		})

	return
}<|endoftext|>


<commit_before>// Copyright 2015 The Mangos Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use 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 respondent implements the RESPONDENT protocol.  This protocol
// receives SURVEYOR requests, and responds with an answer.
package respondent

import (
	"encoding/binary"
	"sync"
	"time"

	"github.com/go-mangos/mangos"
)

type resp struct {
	sock      mangos.ProtocolSocket
	peers     map[uint32]*respPeer
	raw       bool
	ttl       int
	backbuf   []byte
	backtrace []byte
	w         mangos.Waiter
	init      sync.Once
	sync.Mutex
}

type respPeer struct {
	q  chan *mangos.Message
	ep mangos.Endpoint
	x  *resp
}

func (x *resp) Init(sock mangos.ProtocolSocket) {
	x.sock = sock
	x.ttl = 8
	x.peers = make(map[uint32]*respPeer)
	x.w.Init()
	x.backbuf = make([]byte, 0, 64)
	x.sock.SetSendError(mangos.ErrProtoState)
}

func (x *resp) Shutdown(expire time.Time) {
	peers := make(map[uint32]*respPeer)
	x.w.WaitAbsTimeout(expire)
	x.Lock()
	for id, peer := range x.peers {
		delete(x.peers, id)
		peers[id] = peer
	}
	x.Unlock()

	for id, peer := range peers {
		delete(peers, id)
		mangos.DrainChannel(peer.q, expire)
		close(peer.q)
	}
}

func (x *resp) sender() {
	// This is pretty easy because we have only one peer at a time.
	// If the peer goes away, we'll just drop the message on the floor.

	defer x.w.Done()
	cq := x.sock.CloseChannel()
	sq := x.sock.SendChannel()
	for {
		var m *mangos.Message
		select {
		case m = <-sq:
		case <-cq:
			return
		}

		// Lop off the 32-bit peer/pipe ID.  If absent, drop.
		if len(m.Header) < 4 {
			m.Free()
			continue
		}

		id := binary.BigEndian.Uint32(m.Header)
		m.Header = m.Header[4:]

		x.Lock()
		peer := x.peers[id]
		x.Unlock()

		if peer == nil {
			m.Free()
			continue
		}

		// Put it on the outbound queue
		select {
		case peer.q <- m:
		default:
			// Backpressure, drop it.
			m.Free()
		}
	}
}

// When sending, we should have the survey ID in the header.
func (peer *respPeer) sender() {
	for {
		m := <-peer.q
		if m == nil {
			break
		}
		if peer.ep.SendMsg(m) != nil {
			m.Free()
			break
		}
	}
}

func (x *resp) receiver(ep mangos.Endpoint) {

	rq := x.sock.RecvChannel()
	cq := x.sock.CloseChannel()

	for {
		m := ep.RecvMsg()
		if m == nil {
			return
		}

		v := ep.GetID()
		m.Header = append(m.Header,
			byte(v>>24), byte(v>>16), byte(v>>8), byte(v))
		hops := 0

		for {
			if hops >= x.ttl {
				m.Free() // ErrTooManyHops
				continue
			}
			hops++
			if len(m.Body) < 4 {
				m.Free()
				continue
			}
			m.Header = append(m.Header, m.Body[:4]...)
			m.Body = m.Body[4:]
			if m.Header[len(m.Header)-4]&0x80 != 0 {
				break
			}
		}

		select {
		case rq <- m:
		case <-cq:
			m.Free()
			return
		}
	}
}

func (x *resp) RecvHook(m *mangos.Message) bool {
	if x.raw {
		// Raw mode receivers get the message unadulterated.
		return true
	}

	if len(m.Header) < 4 {
		return false
	}

	x.Lock()
	x.backbuf = x.backbuf[0:0] // avoid allocations
	x.backtrace = append(x.backbuf, m.Header...)
	x.Unlock()
	x.sock.SetSendError(nil)
	return true
}

func (x *resp) SendHook(m *mangos.Message) bool {
	if x.raw {
		// Raw mode senders expected to have prepared header already.
		return true
	}
	x.sock.SetSendError(mangos.ErrProtoState)
	x.Lock()
	m.Header = append(m.Header[0:0], x.backtrace...)
	x.backtrace = nil
	x.Unlock()
	if len(m.Header) == 0 {
		return false
	}
	return true
}

func (x *resp) AddEndpoint(ep mangos.Endpoint) {
	x.init.Do(func() {
		x.w.Add()
		go x.sender()
	})
	peer := &respPeer{ep: ep, x: x, q: make(chan *mangos.Message, 1)}

	x.Lock()
	x.peers[ep.GetID()] = peer
	x.Unlock()

	go x.receiver(ep)
	go peer.sender()
}

func (x *resp) RemoveEndpoint(ep mangos.Endpoint) {
	id := ep.GetID()

	x.Lock()
	peer := x.peers[id]
	delete(x.peers, id)
	x.Unlock()

	if peer != nil {
		close(peer.q)
	}
}

func (*resp) Number() uint16 {
	return mangos.ProtoRespondent
}

func (*resp) PeerNumber() uint16 {
	return mangos.ProtoSurveyor
}

func (*resp) Name() string {
	return "respondent"
}

func (*resp) PeerName() string {
	return "surveyor"
}

func (x *resp) SetOption(name string, v interface{}) error {
	var ok bool
	switch name {
	case mangos.OptionRaw:
		if x.raw, ok = v.(bool); !ok {
			return mangos.ErrBadValue
		}
		if x.raw {
			x.sock.SetSendError(nil)
		} else {
			x.sock.SetSendError(mangos.ErrProtoState)
		}
		return nil
	case mangos.OptionTTL:
		if ttl, ok := v.(int); !ok {
			return mangos.ErrBadValue
		} else if ttl < 1 || ttl > 255 {
			return mangos.ErrBadValue
		} else {
			x.ttl = ttl
		}
		return nil
	default:
		return mangos.ErrBadOption
	}
}

func (x *resp) GetOption(name string) (interface{}, error) {
	switch name {
	case mangos.OptionRaw:
		return x.raw, nil
	case mangos.OptionTTL:
		return x.ttl, nil
	default:
		return nil, mangos.ErrBadOption
	}
}

// NewSocket allocates a new Socket using the RESPONDENT protocol.
func NewSocket() (mangos.Socket, error) {
	return mangos.MakeSocket(&resp{}), nil
}
<commit_msg>fixes #177 Benchmark test times out after 10 minutes on OSX<commit_after>// Copyright 2016 The Mangos Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use 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 respondent implements the RESPONDENT protocol.  This protocol
// receives SURVEYOR requests, and responds with an answer.
package respondent

import (
	"encoding/binary"
	"sync"
	"time"

	"github.com/go-mangos/mangos"
)

type resp struct {
	sock      mangos.ProtocolSocket
	peers     map[uint32]*respPeer
	raw       bool
	ttl       int
	backbuf   []byte
	backtrace []byte
	w         mangos.Waiter
	init      sync.Once
	sync.Mutex
}

type respPeer struct {
	q  chan *mangos.Message
	ep mangos.Endpoint
	x  *resp
}

func (x *resp) Init(sock mangos.ProtocolSocket) {
	x.sock = sock
	x.ttl = 8
	x.peers = make(map[uint32]*respPeer)
	x.w.Init()
	x.backbuf = make([]byte, 0, 64)
	x.sock.SetSendError(mangos.ErrProtoState)
}

func (x *resp) Shutdown(expire time.Time) {
	peers := make(map[uint32]*respPeer)
	x.w.WaitAbsTimeout(expire)
	x.Lock()
	for id, peer := range x.peers {
		delete(x.peers, id)
		peers[id] = peer
	}
	x.Unlock()

	for id, peer := range peers {
		delete(peers, id)
		mangos.DrainChannel(peer.q, expire)
		close(peer.q)
	}
}

func (x *resp) sender() {
	// This is pretty easy because we have only one peer at a time.
	// If the peer goes away, we'll just drop the message on the floor.

	defer x.w.Done()
	cq := x.sock.CloseChannel()
	sq := x.sock.SendChannel()
	for {
		var m *mangos.Message
		select {
		case m = <-sq:
		case <-cq:
			return
		}

		// Lop off the 32-bit peer/pipe ID.  If absent, drop.
		if len(m.Header) < 4 {
			m.Free()
			continue
		}

		id := binary.BigEndian.Uint32(m.Header)
		m.Header = m.Header[4:]

		x.Lock()
		peer := x.peers[id]
		x.Unlock()

		if peer == nil {
			m.Free()
			continue
		}

		// Put it on the outbound queue
		select {
		case peer.q <- m:
		default:
			// Backpressure, drop it.
			m.Free()
		}
	}
}

// When sending, we should have the survey ID in the header.
func (peer *respPeer) sender() {
	for {
		m := <-peer.q
		if m == nil {
			break
		}
		if peer.ep.SendMsg(m) != nil {
			m.Free()
			break
		}
	}
}

func (x *resp) receiver(ep mangos.Endpoint) {

	rq := x.sock.RecvChannel()
	cq := x.sock.CloseChannel()

outer:
	for {
		m := ep.RecvMsg()
		if m == nil {
			return
		}

		v := ep.GetID()
		m.Header = append(m.Header,
			byte(v>>24), byte(v>>16), byte(v>>8), byte(v))
		hops := 0

		for {
			if hops >= x.ttl {
				m.Free() // ErrTooManyHops
				continue outer
			}
			hops++
			if len(m.Body) < 4 {
				m.Free()
				continue outer
			}
			m.Header = append(m.Header, m.Body[:4]...)
			m.Body = m.Body[4:]
			if m.Header[len(m.Header)-4]&0x80 != 0 {
				break
			}
		}

		select {
		case rq <- m:
		case <-cq:
			m.Free()
			return
		}
	}
}

func (x *resp) RecvHook(m *mangos.Message) bool {
	if x.raw {
		// Raw mode receivers get the message unadulterated.
		return true
	}

	if len(m.Header) < 4 {
		return false
	}

	x.Lock()
	x.backbuf = x.backbuf[0:0] // avoid allocations
	x.backtrace = append(x.backbuf, m.Header...)
	x.Unlock()
	x.sock.SetSendError(nil)
	return true
}

func (x *resp) SendHook(m *mangos.Message) bool {
	if x.raw {
		// Raw mode senders expected to have prepared header already.
		return true
	}
	x.sock.SetSendError(mangos.ErrProtoState)
	x.Lock()
	m.Header = append(m.Header[0:0], x.backtrace...)
	x.backtrace = nil
	x.Unlock()
	if len(m.Header) == 0 {
		return false
	}
	return true
}

func (x *resp) AddEndpoint(ep mangos.Endpoint) {
	x.init.Do(func() {
		x.w.Add()
		go x.sender()
	})
	peer := &respPeer{ep: ep, x: x, q: make(chan *mangos.Message, 1)}

	x.Lock()
	x.peers[ep.GetID()] = peer
	x.Unlock()

	go x.receiver(ep)
	go peer.sender()
}

func (x *resp) RemoveEndpoint(ep mangos.Endpoint) {
	id := ep.GetID()

	x.Lock()
	peer := x.peers[id]
	delete(x.peers, id)
	x.Unlock()

	if peer != nil {
		close(peer.q)
	}
}

func (*resp) Number() uint16 {
	return mangos.ProtoRespondent
}

func (*resp) PeerNumber() uint16 {
	return mangos.ProtoSurveyor
}

func (*resp) Name() string {
	return "respondent"
}

func (*resp) PeerName() string {
	return "surveyor"
}

func (x *resp) SetOption(name string, v interface{}) error {
	var ok bool
	switch name {
	case mangos.OptionRaw:
		if x.raw, ok = v.(bool); !ok {
			return mangos.ErrBadValue
		}
		if x.raw {
			x.sock.SetSendError(nil)
		} else {
			x.sock.SetSendError(mangos.ErrProtoState)
		}
		return nil
	case mangos.OptionTTL:
		if ttl, ok := v.(int); !ok {
			return mangos.ErrBadValue
		} else if ttl < 1 || ttl > 255 {
			return mangos.ErrBadValue
		} else {
			x.ttl = ttl
		}
		return nil
	default:
		return mangos.ErrBadOption
	}
}

func (x *resp) GetOption(name string) (interface{}, error) {
	switch name {
	case mangos.OptionRaw:
		return x.raw, nil
	case mangos.OptionTTL:
		return x.ttl, nil
	default:
		return nil, mangos.ErrBadOption
	}
}

// NewSocket allocates a new Socket using the RESPONDENT protocol.
func NewSocket() (mangos.Socket, error) {
	return mangos.MakeSocket(&resp{}), nil
}
<|endoftext|>


<commit_before>// Copyright 2016 Google LLC
//
// 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 pubsub

import (
	"context"
	"fmt"
	"testing"
	"time"

	"cloud.google.com/go/internal/testutil"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
	pubsubpb "google.golang.org/genproto/googleapis/pubsub/v1"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

func checkTopicListing(t *testing.T, c *Client, want []string) {
	topics, err := slurpTopics(c.Topics(context.Background()))
	if err != nil {
		t.Fatalf("error listing topics: %v", err)
	}
	var got []string
	for _, topic := range topics {
		got = append(got, topic.ID())
	}
	if !testutil.Equal(got, want) {
		t.Errorf("topic list: got: %v, want: %v", got, want)
	}
}

// All returns the remaining topics from this iterator.
func slurpTopics(it *TopicIterator) ([]*Topic, error) {
	var topics []*Topic
	for {
		switch topic, err := it.Next(); err {
		case nil:
			topics = append(topics, topic)
		case iterator.Done:
			return topics, nil
		default:
			return nil, err
		}
	}
}

func TestTopicID(t *testing.T) {
	const id = "id"
	c, srv := newFake(t)
	defer c.Close()
	defer srv.Close()

	s := c.Topic(id)
	if got, want := s.ID(), id; got != want {
		t.Errorf("Token.ID() = %q; want %q", got, want)
	}
}

func TestListTopics(t *testing.T) {
	c, srv := newFake(t)
	defer c.Close()
	defer srv.Close()

	var ids []string
	for i := 1; i <= 4; i++ {
		id := fmt.Sprintf("t%d", i)
		ids = append(ids, id)
		mustCreateTopic(t, c, id)
	}
	checkTopicListing(t, c, ids)
}

func TestListCompletelyEmptyTopics(t *testing.T) {
	c, srv := newFake(t)
	defer c.Close()
	defer srv.Close()

	checkTopicListing(t, c, nil)
}

func TestStopPublishOrder(t *testing.T) {
	// Check that Stop doesn't panic if called before Publish.
	// Also that Publish after Stop returns the right error.
	ctx := context.Background()
	c := &Client{projectID: "projid"}
	topic := c.Topic("t")
	topic.Stop()
	r := topic.Publish(ctx, &Message{})
	_, err := r.Get(ctx)
	if err != errTopicStopped {
		t.Errorf("got %v, want errTopicStopped", err)
	}
}

func TestPublishTimeout(t *testing.T) {
	ctx := context.Background()
	serv, err := testutil.NewServer()
	if err != nil {
		t.Fatal(err)
	}
	pubsubpb.RegisterPublisherServer(serv.Gsrv, &alwaysFailPublish{})
	conn, err := grpc.Dial(serv.Addr, grpc.WithInsecure())
	if err != nil {
		t.Fatal(err)
	}
	c, err := NewClient(ctx, "projectID", option.WithGRPCConn(conn))
	if err != nil {
		t.Fatal(err)
	}
	topic := c.Topic("t")
	topic.PublishSettings.Timeout = 3 * time.Second
	r := topic.Publish(ctx, &Message{})
	defer topic.Stop()
	select {
	case <-r.Ready():
		_, err = r.Get(ctx)
		if err != context.DeadlineExceeded {
			t.Fatalf("got %v, want context.DeadlineExceeded", err)
		}
	case <-time.After(2 * topic.PublishSettings.Timeout):
		t.Fatal("timed out")
	}
}

func TestUpdateTopic(t *testing.T) {
	ctx := context.Background()
	client, srv := newFake(t)
	defer client.Close()
	defer srv.Close()

	topic := mustCreateTopic(t, client, "T")
	config, err := topic.Config(ctx)
	if err != nil {
		t.Fatal(err)
	}
	want := TopicConfig{}
	if !testutil.Equal(config, want) {
		t.Errorf("got %+v, want %+v", config, want)
	}

	// replace labels
	labels := map[string]string{"label": "value"}
	config2, err := topic.Update(ctx, TopicConfigToUpdate{Labels: labels})
	if err != nil {
		t.Fatal(err)
	}
	want = TopicConfig{
		Labels:               labels,
		MessageStoragePolicy: MessageStoragePolicy{[]string{"US"}},
	}
	if !testutil.Equal(config2, want) {
		t.Errorf("got %+v, want %+v", config2, want)
	}

	// delete all labels
	labels = map[string]string{}
	config3, err := topic.Update(ctx, TopicConfigToUpdate{Labels: labels})
	if err != nil {
		t.Fatal(err)
	}
	want.Labels = nil
	if !testutil.Equal(config3, want) {
		t.Errorf("got %+v, want %+v", config3, want)
	}
}

type alwaysFailPublish struct {
	pubsubpb.PublisherServer
}

func (s *alwaysFailPublish) Publish(ctx context.Context, req *pubsubpb.PublishRequest) (*pubsubpb.PublishResponse, error) {
	return nil, status.Errorf(codes.Unavailable, "try again")
}

func mustCreateTopic(t *testing.T, c *Client, id string) *Topic {
	topic, err := c.CreateTopic(context.Background(), id)
	if err != nil {
		t.Fatal(err)
	}
	return topic
}
<commit_msg>pubsub: fix TestPublishTimeout<commit_after>// Copyright 2016 Google LLC
//
// 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 pubsub

import (
	"context"
	"fmt"
	"testing"
	"time"

	"cloud.google.com/go/internal/testutil"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
	pubsubpb "google.golang.org/genproto/googleapis/pubsub/v1"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

func checkTopicListing(t *testing.T, c *Client, want []string) {
	topics, err := slurpTopics(c.Topics(context.Background()))
	if err != nil {
		t.Fatalf("error listing topics: %v", err)
	}
	var got []string
	for _, topic := range topics {
		got = append(got, topic.ID())
	}
	if !testutil.Equal(got, want) {
		t.Errorf("topic list: got: %v, want: %v", got, want)
	}
}

// All returns the remaining topics from this iterator.
func slurpTopics(it *TopicIterator) ([]*Topic, error) {
	var topics []*Topic
	for {
		switch topic, err := it.Next(); err {
		case nil:
			topics = append(topics, topic)
		case iterator.Done:
			return topics, nil
		default:
			return nil, err
		}
	}
}

func TestTopicID(t *testing.T) {
	const id = "id"
	c, srv := newFake(t)
	defer c.Close()
	defer srv.Close()

	s := c.Topic(id)
	if got, want := s.ID(), id; got != want {
		t.Errorf("Token.ID() = %q; want %q", got, want)
	}
}

func TestListTopics(t *testing.T) {
	c, srv := newFake(t)
	defer c.Close()
	defer srv.Close()

	var ids []string
	for i := 1; i <= 4; i++ {
		id := fmt.Sprintf("t%d", i)
		ids = append(ids, id)
		mustCreateTopic(t, c, id)
	}
	checkTopicListing(t, c, ids)
}

func TestListCompletelyEmptyTopics(t *testing.T) {
	c, srv := newFake(t)
	defer c.Close()
	defer srv.Close()

	checkTopicListing(t, c, nil)
}

func TestStopPublishOrder(t *testing.T) {
	// Check that Stop doesn't panic if called before Publish.
	// Also that Publish after Stop returns the right error.
	ctx := context.Background()
	c := &Client{projectID: "projid"}
	topic := c.Topic("t")
	topic.Stop()
	r := topic.Publish(ctx, &Message{})
	_, err := r.Get(ctx)
	if err != errTopicStopped {
		t.Errorf("got %v, want errTopicStopped", err)
	}
}

func TestPublishTimeout(t *testing.T) {
	ctx := context.Background()
	serv, err := testutil.NewServer()
	if err != nil {
		t.Fatal(err)
	}
	serv.Start()
	pubsubpb.RegisterPublisherServer(serv.Gsrv, &alwaysFailPublish{})
	conn, err := grpc.Dial(serv.Addr, grpc.WithInsecure())
	if err != nil {
		t.Fatal(err)
	}
	c, err := NewClient(ctx, "projectID", option.WithGRPCConn(conn))
	if err != nil {
		t.Fatal(err)
	}
	topic := c.Topic("t")
	topic.PublishSettings.Timeout = 3 * time.Second
	r := topic.Publish(ctx, &Message{})
	defer topic.Stop()
	select {
	case <-r.Ready():
		_, err = r.Get(ctx)
		if err != context.DeadlineExceeded {
			t.Fatalf("got %v, want context.DeadlineExceeded", err)
		}
	case <-time.After(2 * topic.PublishSettings.Timeout):
		t.Fatal("timed out")
	}
}

func TestUpdateTopic(t *testing.T) {
	ctx := context.Background()
	client, srv := newFake(t)
	defer client.Close()
	defer srv.Close()

	topic := mustCreateTopic(t, client, "T")
	config, err := topic.Config(ctx)
	if err != nil {
		t.Fatal(err)
	}
	want := TopicConfig{}
	if !testutil.Equal(config, want) {
		t.Errorf("got %+v, want %+v", config, want)
	}

	// replace labels
	labels := map[string]string{"label": "value"}
	config2, err := topic.Update(ctx, TopicConfigToUpdate{Labels: labels})
	if err != nil {
		t.Fatal(err)
	}
	want = TopicConfig{
		Labels:               labels,
		MessageStoragePolicy: MessageStoragePolicy{[]string{"US"}},
	}
	if !testutil.Equal(config2, want) {
		t.Errorf("got %+v, want %+v", config2, want)
	}

	// delete all labels
	labels = map[string]string{}
	config3, err := topic.Update(ctx, TopicConfigToUpdate{Labels: labels})
	if err != nil {
		t.Fatal(err)
	}
	want.Labels = nil
	if !testutil.Equal(config3, want) {
		t.Errorf("got %+v, want %+v", config3, want)
	}
}

type alwaysFailPublish struct {
	pubsubpb.PublisherServer
}

func (s *alwaysFailPublish) Publish(ctx context.Context, req *pubsubpb.PublishRequest) (*pubsubpb.PublishResponse, error) {
	return nil, status.Errorf(codes.Unavailable, "try again")
}

func mustCreateTopic(t *testing.T, c *Client, id string) *Topic {
	topic, err := c.CreateTopic(context.Background(), id)
	if err != nil {
		t.Fatal(err)
	}
	return topic
}
<|endoftext|>


<commit_before>package pushbullet

import (
	"bytes"
	"encoding/json"
	"errors"
	"net/http"

	"github.com/mitsuse/bullet/pushbullet/requests"
)

// Push a note, which consists of "title" and "message" strings.
func (pb *Pushbullet) PostPushesNote(n *requests.Note) error {
	return pb.postPushes(n)
}

// Push a link, which consists of "title", "message" and "url" strings.
func (pb *Pushbullet) PostPushesLink(l *requests.Link) error {
	return pb.postPushes(l)
}

// Push an address, which consists of the place "name" and "address (searchquery)" for map.
func (pb *Pushbullet) PostPushesAddress(a *requests.Address) error {
	return pb.postPushes(a)
}

// Push a checklist, which consists of "title" and the list of items.
func (pb *Pushbullet) PostPushesChecklist(c *requests.Checklist) error {
	return pb.postPushes(c)
}

// Push a file, which consists of "title", "message" and the information of uploaded file.
func (pb *Pushbullet) PostPushesFile(f *requests.File) error {
	return pb.postPushes(f)
}

// func (pb *Pushbullet) postPushes(p pushes.Push) error {
func (pb *Pushbullet) postPushes(p interface{}) error {
	buffer := &bytes.Buffer{}

	encoder := json.NewEncoder(buffer)
	if err := encoder.Encode(p); err != nil {
		return err
	}

	req, err := http.NewRequest("POST", ENDPOINT_PUSHES, buffer)
	if err != nil {
		return err
	}

	req.Header.Add("Content-Type", "application/json")
	req.SetBasicAuth(pb.token, "")

	// TODO: Set the timeout.
	client := &http.Client{}

	res, err := client.Do(req)
	if err != nil {
		return err
	}

	// TODO: Return an error value with human friendly message.
	if res.StatusCode != 200 {
		return errors.New(res.Status)
	}

	return nil
}
<commit_msg>Remove a needless comment.<commit_after>package pushbullet

import (
	"bytes"
	"encoding/json"
	"errors"
	"net/http"

	"github.com/mitsuse/bullet/pushbullet/requests"
)

// Push a note, which consists of "title" and "message" strings.
func (pb *Pushbullet) PostPushesNote(n *requests.Note) error {
	return pb.postPushes(n)
}

// Push a link, which consists of "title", "message" and "url" strings.
func (pb *Pushbullet) PostPushesLink(l *requests.Link) error {
	return pb.postPushes(l)
}

// Push an address, which consists of the place "name" and "address (searchquery)" for map.
func (pb *Pushbullet) PostPushesAddress(a *requests.Address) error {
	return pb.postPushes(a)
}

// Push a checklist, which consists of "title" and the list of items.
func (pb *Pushbullet) PostPushesChecklist(c *requests.Checklist) error {
	return pb.postPushes(c)
}

// Push a file, which consists of "title", "message" and the information of uploaded file.
func (pb *Pushbullet) PostPushesFile(f *requests.File) error {
	return pb.postPushes(f)
}

func (pb *Pushbullet) postPushes(p interface{}) error {
	buffer := &bytes.Buffer{}

	encoder := json.NewEncoder(buffer)
	if err := encoder.Encode(p); err != nil {
		return err
	}

	req, err := http.NewRequest("POST", ENDPOINT_PUSHES, buffer)
	if err != nil {
		return err
	}

	req.Header.Add("Content-Type", "application/json")
	req.SetBasicAuth(pb.token, "")

	// TODO: Set the timeout.
	client := &http.Client{}

	res, err := client.Do(req)
	if err != nil {
		return err
	}

	// TODO: Return an error value with human friendly message.
	if res.StatusCode != 200 {
		return errors.New(res.Status)
	}

	return nil
}
<|endoftext|>


<commit_before>package commandlogger

import (
	log "github.com/sirupsen/logrus"
	pb "gitlab.com/gitlab-org/gitaly/proto/go/gitalypb"
	"gitlab.com/gitlab-org/gitlab-shell/internal/gitlabnet/accessverifier"
)

func Log(command string, repository *pb.Repository, response *accessverifier.Response, protocol string) {
	fields := log.Fields{
		"command":       command,
		"glProjectPath": repository.GlProjectPath,
		"glRepository":  repository.GlRepository,
		"userId":        response.UserId,
		"userName":      response.Username,
		"gitProtocol":   protocol,
	}

	log.WithFields(fields).Info("executing git command")
}
<commit_msg>Change git command logging keys to be be snake cased<commit_after>package commandlogger

import (
	log "github.com/sirupsen/logrus"
	pb "gitlab.com/gitlab-org/gitaly/proto/go/gitalypb"
	"gitlab.com/gitlab-org/gitlab-shell/internal/gitlabnet/accessverifier"
)

func Log(command string, repository *pb.Repository, response *accessverifier.Response, protocol string) {
	fields := log.Fields{
		"command":       command,
		"gl_project_path": repository.GlProjectPath,
		"gl_repository":  repository.GlRepository,
		"user_id":        response.UserId,
		"username":      response.Username,
		"git_protocol":   protocol,
	}

	log.WithFields(fields).Info("executing git command")
}
<|endoftext|>


<commit_before>package drivers

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

	log "gopkg.in/inconshreveable/log15.v2"

	"github.com/lxc/lxd/lxd/instance/drivers/qmp"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/metrics"
	"github.com/lxc/lxd/shared"
)

func (d *qemu) getQemuMetrics() (*metrics.MetricSet, error) {
	// Connect to the monitor.
	monitor, err := qmp.Connect(d.monitorPath(), qemuSerialChardevName, d.getMonitorEventHandler())
	if err != nil {
		return nil, err
	}

	out := metrics.Metrics{}

	cpuStats, err := d.getQemuCPUMetrics(monitor)
	if err != nil {
		d.logger.Warn("Failed to get CPU metrics", log.Ctx{"err": err})
	} else {
		out.CPU = cpuStats
	}

	memoryStats, err := d.getQemuMemoryMetrics(monitor)
	if err != nil {
		d.logger.Warn("Failed to get memory metrics", log.Ctx{"err": err})
	} else {
		out.Memory = memoryStats
	}

	diskStats, err := d.getQemuDiskMetrics(monitor)
	if err != nil {
		d.logger.Warn("Failed to get memory metrics", log.Ctx{"err": err})
	} else {
		out.Disk = diskStats
	}

	networkState, err := d.getNetworkState()
	if err != nil {
		d.logger.Warn("Failed to get network metrics", log.Ctx{"err": err})
	} else {

		out.Network = make(map[string]metrics.NetworkMetrics)

		for name, state := range networkState {
			out.Network[name] = metrics.NetworkMetrics{
				ReceiveBytes:    uint64(state.Counters.BytesReceived),
				ReceiveDrop:     uint64(state.Counters.PacketsDroppedInbound),
				ReceiveErrors:   uint64(state.Counters.ErrorsReceived),
				ReceivePackets:  uint64(state.Counters.PacketsReceived),
				TransmitBytes:   uint64(state.Counters.BytesSent),
				TransmitDrop:    uint64(state.Counters.PacketsDroppedOutbound),
				TransmitErrors:  uint64(state.Counters.ErrorsSent),
				TransmitPackets: uint64(state.Counters.PacketsSent),
			}
		}
	}

	metricSet, err := metrics.MetricSetFromAPI(&out, map[string]string{"project": d.project, "name": d.name, "type": instancetype.VM.String()})
	if err != nil {
		return nil, err
	}

	return metricSet, nil
}

func (d *qemu) getQemuDiskMetrics(monitor *qmp.Monitor) (map[string]metrics.DiskMetrics, error) {
	stats, err := monitor.GetBlockStats()
	if err != nil {
		return nil, err
	}

	out := make(map[string]metrics.DiskMetrics)

	for dev, stat := range stats {
		out[dev] = metrics.DiskMetrics{
			ReadBytes:       uint64(stat.BytesRead),
			ReadsCompleted:  uint64(stat.ReadsCompleted),
			WrittenBytes:    uint64(stat.BytesWritten),
			WritesCompleted: uint64(stat.WritesCompleted),
		}
	}

	return out, nil
}

func (d *qemu) getQemuMemoryMetrics(monitor *qmp.Monitor) (metrics.MemoryMetrics, error) {
	stats, err := monitor.GetMemoryStats()
	if err != nil {
		return metrics.MemoryMetrics{}, err
	}

	out := metrics.MemoryMetrics{
		MemAvailableBytes: uint64(stats.AvailableMemory),
		MemFreeBytes:      uint64(stats.FreeMemory),
		MemTotalBytes:     uint64(stats.TotalMemory),
	}

	return out, nil
}

func (d *qemu) getQemuCPUMetrics(monitor *qmp.Monitor) (map[string]metrics.CPUMetrics, error) {
	// Get CPU metrics
	threadIDs, err := monitor.GetCPUs()
	if err != nil {
		return nil, err
	}

	cpuMetrics := map[string]metrics.CPUMetrics{}

	for i, threadID := range threadIDs {
		pid, err := ioutil.ReadFile(d.pidFilePath())
		if err != nil {
			return nil, err
		}

		statFile := filepath.Join("/proc", strings.TrimSpace(string(pid)), "task", strconv.Itoa(threadID), "stat")

		if !shared.PathExists(statFile) {
			continue
		}

		content, err := ioutil.ReadFile(statFile)
		if err != nil {
			return nil, err
		}

		fields := strings.Fields(string(content))

		stats := metrics.CPUMetrics{}

		stats.SecondsUser, err = strconv.ParseFloat(fields[13], 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[13], err)
		}

		guestTime, err := strconv.ParseFloat(fields[42], 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[42], err)
		}

		// According to proc(5), utime includes guest_time which therefore needs to be subtracted to get the correct time.
		stats.SecondsUser -= guestTime
		stats.SecondsUser /= 100

		stats.SecondsSystem, err = strconv.ParseFloat(fields[14], 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[14], err)
		}

		stats.SecondsSystem /= 100

		cpuMetrics[fmt.Sprintf("cpu%d", i)] = stats
	}

	return cpuMetrics, nil
}
<commit_msg>lxd/instance/qemu: Fix agent-less memory metrics<commit_after>package drivers

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

	log "gopkg.in/inconshreveable/log15.v2"

	"github.com/lxc/lxd/lxd/instance/drivers/qmp"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/metrics"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/units"
)

func (d *qemu) getQemuMetrics() (*metrics.MetricSet, error) {
	// Connect to the monitor.
	monitor, err := qmp.Connect(d.monitorPath(), qemuSerialChardevName, d.getMonitorEventHandler())
	if err != nil {
		return nil, err
	}

	out := metrics.Metrics{}

	cpuStats, err := d.getQemuCPUMetrics(monitor)
	if err != nil {
		d.logger.Warn("Failed to get CPU metrics", log.Ctx{"err": err})
	} else {
		out.CPU = cpuStats
	}

	memoryStats, err := d.getQemuMemoryMetrics(monitor)
	if err != nil {
		d.logger.Warn("Failed to get memory metrics", log.Ctx{"err": err})
	} else {
		out.Memory = memoryStats
	}

	diskStats, err := d.getQemuDiskMetrics(monitor)
	if err != nil {
		d.logger.Warn("Failed to get disk metrics", log.Ctx{"err": err})
	} else {
		out.Disk = diskStats
	}

	networkState, err := d.getNetworkState()
	if err != nil {
		d.logger.Warn("Failed to get network metrics", log.Ctx{"err": err})
	} else {

		out.Network = make(map[string]metrics.NetworkMetrics)

		for name, state := range networkState {
			out.Network[name] = metrics.NetworkMetrics{
				ReceiveBytes:    uint64(state.Counters.BytesReceived),
				ReceiveDrop:     uint64(state.Counters.PacketsDroppedInbound),
				ReceiveErrors:   uint64(state.Counters.ErrorsReceived),
				ReceivePackets:  uint64(state.Counters.PacketsReceived),
				TransmitBytes:   uint64(state.Counters.BytesSent),
				TransmitDrop:    uint64(state.Counters.PacketsDroppedOutbound),
				TransmitErrors:  uint64(state.Counters.ErrorsSent),
				TransmitPackets: uint64(state.Counters.PacketsSent),
			}
		}
	}

	metricSet, err := metrics.MetricSetFromAPI(&out, map[string]string{"project": d.project, "name": d.name, "type": instancetype.VM.String()})
	if err != nil {
		return nil, err
	}

	return metricSet, nil
}

func (d *qemu) getQemuDiskMetrics(monitor *qmp.Monitor) (map[string]metrics.DiskMetrics, error) {
	stats, err := monitor.GetBlockStats()
	if err != nil {
		return nil, err
	}

	out := make(map[string]metrics.DiskMetrics)

	for dev, stat := range stats {
		out[dev] = metrics.DiskMetrics{
			ReadBytes:       uint64(stat.BytesRead),
			ReadsCompleted:  uint64(stat.ReadsCompleted),
			WrittenBytes:    uint64(stat.BytesWritten),
			WritesCompleted: uint64(stat.WritesCompleted),
		}
	}

	return out, nil
}

func (d *qemu) getQemuMemoryMetrics(monitor *qmp.Monitor) (metrics.MemoryMetrics, error) {
	out := metrics.MemoryMetrics{}

	// Get the QEMU PID.
	pid, err := d.pid()
	if err != nil {
		return out, err
	}

	// Extract current QEMU RSS.
	f, err := os.Open(fmt.Sprintf("/proc/%d/status", pid))
	if err != nil {
		return out, err
	}
	defer f.Close()

	// Read it line by line.
	memRSS := int64(-1)

	scan := bufio.NewScanner(f)
	for scan.Scan() {
		line := scan.Text()

		// We only care about VmRSS.
		if !strings.HasPrefix(line, "VmRSS:") {
			continue
		}

		// Extract the before last (value) and last (unit) fields
		fields := strings.Split(line, "\t")
		value := strings.Replace(fields[len(fields)-1], " ", "", -1)

		// Feed the result to units.ParseByteSizeString to get an int value
		valueBytes, err := units.ParseByteSizeString(value)
		if err != nil {
			return out, err
		}

		memRSS = valueBytes
		break
	}

	if memRSS == -1 {
		return out, fmt.Errorf("Couldn't find VM memory usage")
	}

	// Get max memory usage.
	memTotal := d.expandedConfig["limits.memory"]
	if memTotal == "" {
		memTotal = qemuDefaultMemSize // Default if no memory limit specified.
	}

	memTotalBytes, err := units.ParseByteSizeString(memTotal)
	if err != nil {
		return out, err
	}

	// Prepare struct.
	out = metrics.MemoryMetrics{
		MemAvailableBytes: uint64(memTotalBytes - memRSS),
		MemFreeBytes:      uint64(memTotalBytes - memRSS),
		MemTotalBytes:     uint64(memTotalBytes),
	}

	return out, nil
}

func (d *qemu) getQemuCPUMetrics(monitor *qmp.Monitor) (map[string]metrics.CPUMetrics, error) {
	// Get CPU metrics
	threadIDs, err := monitor.GetCPUs()
	if err != nil {
		return nil, err
	}

	cpuMetrics := map[string]metrics.CPUMetrics{}

	for i, threadID := range threadIDs {
		pid, err := ioutil.ReadFile(d.pidFilePath())
		if err != nil {
			return nil, err
		}

		statFile := filepath.Join("/proc", strings.TrimSpace(string(pid)), "task", strconv.Itoa(threadID), "stat")

		if !shared.PathExists(statFile) {
			continue
		}

		content, err := ioutil.ReadFile(statFile)
		if err != nil {
			return nil, err
		}

		fields := strings.Fields(string(content))

		stats := metrics.CPUMetrics{}

		stats.SecondsUser, err = strconv.ParseFloat(fields[13], 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[13], err)
		}

		guestTime, err := strconv.ParseFloat(fields[42], 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[42], err)
		}

		// According to proc(5), utime includes guest_time which therefore needs to be subtracted to get the correct time.
		stats.SecondsUser -= guestTime
		stats.SecondsUser /= 100

		stats.SecondsSystem, err = strconv.ParseFloat(fields[14], 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[14], err)
		}

		stats.SecondsSystem /= 100

		cpuMetrics[fmt.Sprintf("cpu%d", i)] = stats
	}

	return cpuMetrics, nil
}
<|endoftext|>


<commit_before>package containerd

import (
	"context"
	"io"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/images/archive"
	"github.com/containerd/containerd/platforms"
	"github.com/docker/distribution/reference"
)

// ExportImage exports a list of images to the given output stream. The
// exported images are archived into a tar when written to the output
// stream. All images with the given tag and all versions containing
// the same tag are exported. names is the set of tags to export, and
// outStream is the writer which the images are written to.
func (i *ImageService) ExportImage(ctx context.Context, names []string, outStream io.Writer) error {
	opts := []archive.ExportOpt{
		archive.WithPlatform(platforms.Ordered(platforms.DefaultSpec())),
		archive.WithSkipNonDistributableBlobs(),
	}
	is := i.client.ImageService()
	for _, imageRef := range names {
		named, err := reference.ParseDockerRef(imageRef)
		if err != nil {
			return err
		}
		opts = append(opts, archive.WithImage(is, named.String()))
	}
	return i.client.Export(ctx, outStream, opts...)
}

// LoadImage uploads a set of images into the repository. This is the
// complement of ExportImage.  The input stream is an uncompressed tar
// ball containing images and metadata.
func (i *ImageService) LoadImage(ctx context.Context, inTar io.ReadCloser, outStream io.Writer, quiet bool) error {
	_, err := i.client.Import(ctx, inTar,
		containerd.WithImportPlatform(platforms.DefaultStrict()),
	)
	return err
}
<commit_msg>containerd: Unpack loaded images<commit_after>package containerd

import (
	"context"
	"io"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/images/archive"
	"github.com/containerd/containerd/platforms"
	"github.com/docker/distribution/reference"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// ExportImage exports a list of images to the given output stream. The
// exported images are archived into a tar when written to the output
// stream. All images with the given tag and all versions containing
// the same tag are exported. names is the set of tags to export, and
// outStream is the writer which the images are written to.
//
// TODO(thaJeztah): produce JSON stream progress response and image events; see https://github.com/moby/moby/issues/43910
func (i *ImageService) ExportImage(ctx context.Context, names []string, outStream io.Writer) error {
	opts := []archive.ExportOpt{
		archive.WithPlatform(platforms.Ordered(platforms.DefaultSpec())),
		archive.WithSkipNonDistributableBlobs(),
	}
	is := i.client.ImageService()
	for _, imageRef := range names {
		named, err := reference.ParseDockerRef(imageRef)
		if err != nil {
			return err
		}
		opts = append(opts, archive.WithImage(is, named.String()))
	}
	return i.client.Export(ctx, outStream, opts...)
}

// LoadImage uploads a set of images into the repository. This is the
// complement of ExportImage.  The input stream is an uncompressed tar
// ball containing images and metadata.
//
// TODO(thaJeztah): produce JSON stream progress response and image events; see https://github.com/moby/moby/issues/43910
func (i *ImageService) LoadImage(ctx context.Context, inTar io.ReadCloser, outStream io.Writer, quiet bool) error {
	platform := platforms.DefaultStrict()
	imgs, err := i.client.Import(ctx, inTar, containerd.WithImportPlatform(platform))

	if err != nil {
		// TODO(thaJeztah): remove this log or change to debug once we can; see https://github.com/moby/moby/pull/43822#discussion_r937502405
		logrus.WithError(err).Warn("failed to import image to containerd")
		return errors.Wrap(err, "failed to import image")
	}

	for _, img := range imgs {
		platformImg := containerd.NewImageWithPlatform(i.client, img, platform)

		unpacked, err := platformImg.IsUnpacked(ctx, containerd.DefaultSnapshotter)
		if err != nil {
			// TODO(thaJeztah): remove this log or change to debug once we can; see https://github.com/moby/moby/pull/43822#discussion_r937502405
			logrus.WithError(err).WithField("image", img.Name).Debug("failed to check if image is unpacked")
			continue
		}

		if !unpacked {
			err := platformImg.Unpack(ctx, containerd.DefaultSnapshotter)
			if err != nil {
				// TODO(thaJeztah): remove this log or change to debug once we can; see https://github.com/moby/moby/pull/43822#discussion_r937502405
				logrus.WithError(err).WithField("image", img.Name).Warn("failed to unpack image")
				return errors.Wrap(err, "failed to unpack image")
			}
		}
	}
	return nil
}
<|endoftext|>