<commit_before>package example

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/golang/protobuf/proto"
	"github.com/google/go-cmp/cmp"
)

func TestHelloWorldServer_SayHello(t *testing.T) {

	var tests = []struct {
		name    string
		reqFunc func() (*http.Request, error)
		cb      func(ctx context.Context, w http.ResponseWriter, r *http.Request, arg, ret proto.Message, err error)
		wantErr bool
		want    *HelloReply
	}{
		{
			name: "Content-Type JSON",
			reqFunc: func() (*http.Request, error) {
				p := &HelloRequest{
					Name: "John",
				}

				body := &bytes.Buffer{}
				if err := json.NewEncoder(body).Encode(p); err != nil {
					return nil, err
				}

				t.Log(body.String())

				req := httptest.NewRequest(http.MethodPost, "/", body)
				req.Header.Set("Content-Type", "application/json")
				return req, nil
			},
			cb:      nil,
			wantErr: false,
			want: &HelloReply{
				Message: "Hello, John!",
			},
		},
	}

	handler := NewGreeterHandler(&HelloWorldServer{})

	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			req, err := tt.reqFunc()
			if err != nil {
				t.Fatal(err)
			}

			rec := httptest.NewRecorder()
			handler.SayHello(tt.cb).ServeHTTP(rec, req)

			resp := &HelloReply{}
			switch req.Header.Get("Content-Type") {
			case "application/protobuf":
				if err := proto.Unmarshal(rec.Body.Bytes(), resp); err != nil {
					t.Fatal(err)
				}
			case "application/json":
				if err := json.NewDecoder(rec.Body).Decode(resp); err != nil {
					t.Fatal(err)
				}
			default:
			}
			if diff := cmp.Diff(resp, tt.want); diff != "" {
				t.Errorf("%s", diff)
			}
		})
	}
}
<commit_msg>Add protobuf test case<commit_after>package example

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/golang/protobuf/proto"
	"github.com/google/go-cmp/cmp"
)

func TestHelloWorldServer_SayHello(t *testing.T) {

	var tests = []struct {
		name    string
		reqFunc func() (*http.Request, error)
		cb      func(ctx context.Context, w http.ResponseWriter, r *http.Request, arg, ret proto.Message, err error)
		wantErr bool
		want    *HelloReply
	}{
		{
			name: "Content-Type JSON",
			reqFunc: func() (*http.Request, error) {
				p := &HelloRequest{
					Name: "John",
				}

				body := &bytes.Buffer{}
				if err := json.NewEncoder(body).Encode(p); err != nil {
					return nil, err
				}

				req := httptest.NewRequest(http.MethodPost, "/", body)
				req.Header.Set("Content-Type", "application/json")
				return req, nil
			},
			cb:      nil,
			wantErr: false,
			want: &HelloReply{
				Message: "Hello, John!",
			},
		},
		{
			name: "Content-Type Protobuf",
			reqFunc: func() (*http.Request, error) {
				p := &HelloRequest{
					Name: "Smith",
				}

				buf, err := proto.Marshal(p)
				if err != nil {
					return nil, err
				}
				body := bytes.NewBuffer(buf)

				req := httptest.NewRequest(http.MethodPost, "/", body)
				req.Header.Set("Content-Type", "application/protobuf")
				return req, nil
			},
			cb:      nil,
			wantErr: false,
			want: &HelloReply{
				Message: "Hello, Smith!",
			},
		},
	}

	handler := NewGreeterHandler(&HelloWorldServer{})

	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			req, err := tt.reqFunc()
			if err != nil {
				t.Fatal(err)
			}

			rec := httptest.NewRecorder()
			handler.SayHello(tt.cb).ServeHTTP(rec, req)

			resp := &HelloReply{}
			switch req.Header.Get("Content-Type") {
			case "application/protobuf":
				if err := proto.Unmarshal(rec.Body.Bytes(), resp); err != nil {
					t.Fatal(err)
				}
			case "application/json":
				if err := json.NewDecoder(rec.Body).Decode(resp); err != nil {
					t.Fatal(err)
				}
			default:
			}
			if diff := cmp.Diff(resp, tt.want); diff != "" {
				t.Errorf("%s", diff)
			}
		})
	}
}
<|endoftext|>


<commit_before>package expr

import (
	"bytes"
	"errors"
	"sort"

	"github.com/grafana/metrictank/api/models"
	"github.com/raintank/schema"
)

type FuncGroupByTags struct {
	in         GraphiteFunc
	aggregator string
	tags       []string
}

func NewGroupByTags() GraphiteFunc {
	return &FuncGroupByTags{}
}

func (s *FuncGroupByTags) Signature() ([]Arg, []Arg) {
	return []Arg{
		ArgSeriesList{val: &s.in},
		ArgString{val: &s.aggregator, validator: []Validator{IsAggFunc}},
		ArgStrings{val: &s.tags},
	}, []Arg{ArgSeries{}}
}

func (s *FuncGroupByTags) Context(context Context) Context {
	return context
}

func (s *FuncGroupByTags) Exec(cache map[Req][]models.Series) ([]models.Series, error) {
	series, err := s.in.Exec(cache)
	if err != nil {
		return nil, err
	}

	if len(series) == 0 {
		return series, nil
	}

	if len(s.tags) == 0 {
		return nil, errors.New("No tags specified")
	}

	groups := make(map[string][]models.Series)
	useName := false

	groupTags := s.tags
	for i, tag := range groupTags {
		if tag == "name" {
			// We handle name explicitly, remove it from tags
			useName = true
			groupTags = append(groupTags[:i], groupTags[i+1:]...)
			break
		}
	}

	nameReplace := ""
	if !useName {
		// if all series have the same name, name becomes one of our tags
		for _, serie := range series {
			thisName := serie.Tags["name"]
			if nameReplace == "" {
				nameReplace = thisName
			} else if nameReplace != thisName {
				nameReplace = s.aggregator
				break
			}
		}
	}

	// Tags need to be sorted
	sort.Strings(groupTags)

	// First pass - group our series together by key
	var buffer bytes.Buffer
	for _, serie := range series {
		buffer.Reset()

		if useName {
			buffer.WriteString(serie.Tags["name"])
		} else {
			buffer.WriteString(nameReplace)
		}

		for _, goal := range groupTags {
			buffer.WriteRune(';')
			buffer.WriteString(goal)
			buffer.WriteRune('=')

			tagVal, ok := serie.Tags[goal]
			if ok {
				buffer.WriteString(tagVal)
			}
		}

		key := buffer.String()

		groups[key] = append(groups[key], serie)
	}

	output := make([]models.Series, 0, len(groups))
	aggFunc := getCrossSeriesAggFunc(s.aggregator)

	// Now, for each key perform the requested aggregation
	cons, queryCons := summarizeCons(series)

	for name, groupSeries := range groups {
		newSeries := models.Series{
			Target:       name,
			QueryPatt:    name,
			Interval:     series[0].Interval,
			Consolidator: cons,
			QueryCons:    queryCons,
		}
		newSeries.SetTags()

		newSeries.Datapoints = pointSlicePool.Get().([]schema.Point)

		aggFunc(groupSeries, &newSeries.Datapoints)
		cache[Req{}] = append(cache[Req{}], newSeries)

		output = append(output, newSeries)
	}

	return output, nil
}
<commit_msg>Call summarizeCons for each group<commit_after>package expr

import (
	"bytes"
	"errors"
	"sort"

	"github.com/grafana/metrictank/api/models"
	"github.com/raintank/schema"
)

type FuncGroupByTags struct {
	in         GraphiteFunc
	aggregator string
	tags       []string
}

func NewGroupByTags() GraphiteFunc {
	return &FuncGroupByTags{}
}

func (s *FuncGroupByTags) Signature() ([]Arg, []Arg) {
	return []Arg{
		ArgSeriesList{val: &s.in},
		ArgString{val: &s.aggregator, validator: []Validator{IsAggFunc}},
		ArgStrings{val: &s.tags},
	}, []Arg{ArgSeries{}}
}

func (s *FuncGroupByTags) Context(context Context) Context {
	return context
}

func (s *FuncGroupByTags) Exec(cache map[Req][]models.Series) ([]models.Series, error) {
	series, err := s.in.Exec(cache)
	if err != nil {
		return nil, err
	}

	if len(series) == 0 {
		return series, nil
	}

	if len(s.tags) == 0 {
		return nil, errors.New("No tags specified")
	}

	groups := make(map[string][]models.Series)
	useName := false

	groupTags := s.tags
	for i, tag := range groupTags {
		if tag == "name" {
			// We handle name explicitly, remove it from tags
			useName = true
			groupTags = append(groupTags[:i], groupTags[i+1:]...)
			break
		}
	}

	nameReplace := ""
	if !useName {
		// if all series have the same name, name becomes one of our tags
		for _, serie := range series {
			thisName := serie.Tags["name"]
			if nameReplace == "" {
				nameReplace = thisName
			} else if nameReplace != thisName {
				nameReplace = s.aggregator
				break
			}
		}
	}

	// Tags need to be sorted
	sort.Strings(groupTags)

	// First pass - group our series together by key
	var buffer bytes.Buffer
	for _, serie := range series {
		buffer.Reset()

		if useName {
			buffer.WriteString(serie.Tags["name"])
		} else {
			buffer.WriteString(nameReplace)
		}

		for _, goal := range groupTags {
			buffer.WriteRune(';')
			buffer.WriteString(goal)
			buffer.WriteRune('=')

			tagVal, ok := serie.Tags[goal]
			if ok {
				buffer.WriteString(tagVal)
			}
		}

		key := buffer.String()

		groups[key] = append(groups[key], serie)
	}

	output := make([]models.Series, 0, len(groups))
	aggFunc := getCrossSeriesAggFunc(s.aggregator)

	// Now, for each key perform the requested aggregation
	for name, groupSeries := range groups {
		cons, queryCons := summarizeCons(groupSeries)

		newSeries := models.Series{
			Target:       name,
			QueryPatt:    name,
			Interval:     series[0].Interval,
			Consolidator: cons,
			QueryCons:    queryCons,
		}
		newSeries.SetTags()

		newSeries.Datapoints = pointSlicePool.Get().([]schema.Point)

		aggFunc(groupSeries, &newSeries.Datapoints)
		cache[Req{}] = append(cache[Req{}], newSeries)

		output = append(output, newSeries)
	}

	return output, nil
}
<|endoftext|>


<commit_before>package expr

import (
	"math"
	"strconv"
	"testing"

	"github.com/grafana/metrictank/api/models"
	"github.com/grafana/metrictank/schema"
	"github.com/grafana/metrictank/test"
)

var datapoints = []schema.Point{
	{Val: 0, Ts: 10},
	{Val: math.NaN(), Ts: 20},
	{Val: -10, Ts: 30},
	{Val: 5.5, Ts: 40},
	{Val: -math.MaxFloat64, Ts: 50},
	{Val: -1234567890, Ts: 60},
	{Val: math.MaxFloat64, Ts: 70},
}

var datapointsInvert = []schema.Point{
	{Val: math.NaN(), Ts: 10},
	{Val: math.NaN(), Ts: 20},
	{Val: math.Pow(-10, -1), Ts: 30},
	{Val: math.Pow(5.5, -1), Ts: 40},
	{Val: math.Pow(-math.MaxFloat64, -1), Ts: 50},
	{Val: math.Pow(-1234567890, -1), Ts: 60},
	{Val: math.Pow(math.MaxFloat64, -1), Ts: 70},
}

func TestInvertBasic(t *testing.T) {
	in := []models.Series{
		{
			Interval:   10,
			QueryPatt:  "queryPattHere",
			Target:     "targetHere",
			Datapoints: getCopy(datapoints),
		},
	}

	f := getNewInvert(in)
	out := []models.Series{
		{
			Interval:   10,
			QueryPatt:  "invert(queryPattHere)",
			Target:     "invert(targetHere)",
			Datapoints: getCopy(datapointsInvert),
		},
	}

	got, err := f.Exec(make(map[Req][]models.Series))
	if err := equalOutput(out, got, nil, err); err != nil {
		t.Fatal(err)
	}
}

func TestInvertInputUnchanged(t *testing.T) {
	in := []models.Series{
		{
			Interval:   10,
			QueryPatt:  "queryPattHere",
			Target:     "targetHere",
			Datapoints: getCopy(datapoints),
		},
	}

	// store copy of the original input
	inCopy := []models.Series{
		{
			Interval:   10,
			QueryPatt:  "queryPattHere",
			Target:     "targetHere",
			Datapoints: getCopy(datapoints),
		},
	}

	f := getNewInvert(in)
	_, err := f.Exec(make(map[Req][]models.Series))

	// make sure input hasn't changed after call to invert
	if err := equalOutput(in, inCopy, nil, err); err != nil {
		t.Fatal(err)
	}
}

func TestInvertEmpty(t *testing.T) {
	in := []models.Series{
		{
			Interval:   10,
			QueryPatt:  "queryPattHere",
			Target:     "targetHere",
			Datapoints: []schema.Point{},
		},
	}

	f := getNewInvert(in)

	out := []models.Series{
		{
			Interval:   10,
			QueryPatt:  "invert(queryPattHere)",
			Target:     "invert(targetHere)",
			Datapoints: []schema.Point{},
		},
	}

	got, err := f.Exec(make(map[Req][]models.Series))
	if err := equalOutput(out, got, nil, err); err != nil {
		t.Fatal(err)
	}
}

func TestInvertMultipleSeries(t *testing.T) {
	in := []models.Series{
		{
			Interval:   10,
			QueryPatt:  "queryPattHere",
			Target:     "targetHere",
			Datapoints: getCopy(datapoints),
		},
		{
			Interval:   20,
			QueryPatt:  "queryPattHere2",
			Target:     "targetHere2",
			Datapoints: getCopy(datapoints),
		},
	}

	f := getNewInvert(in)
	out := []models.Series{
		{
			Interval:   10,
			QueryPatt:  "invert(queryPattHere)",
			Target:     "invert(targetHere)",
			Datapoints: getCopy(datapointsInvert),
		},
		{
			Interval:   20,
			QueryPatt:  "invert(queryPattHere2)",
			Target:     "invert(targetHere2)",
			Datapoints: getCopy(datapointsInvert),
		},
	}

	got, err := f.Exec(make(map[Req][]models.Series))
	if err := equalOutput(out, got, nil, err); err != nil {
		t.Fatal(err)
	}
}

func getNewInvert(in []models.Series) *FuncInvert {
	f := NewInvert()
	ps := f.(*FuncInvert)
	ps.in = NewMock(in)
	return ps
}

func BenchmarkInvert10k_1NoNulls(b *testing.B) {
	benchmarkInvert(b, 1, test.RandFloats10k, test.RandFloats10k)
}
func BenchmarkInvert10k_10NoNulls(b *testing.B) {
	benchmarkInvert(b, 10, test.RandFloats10k, test.RandFloats10k)
}
func BenchmarkInvert10k_100NoNulls(b *testing.B) {
	benchmarkInvert(b, 100, test.RandFloats10k, test.RandFloats10k)
}
func BenchmarkInvert10k_1000NoNulls(b *testing.B) {
	benchmarkInvert(b, 1000, test.RandFloats10k, test.RandFloats10k)
}
func BenchmarkInvert10k_1SomeSeriesHalfNulls(b *testing.B) {
	benchmarkInvert(b, 1, test.RandFloats10k, test.RandFloatsWithNulls10k)
}
func BenchmarkInvert10k_10SomeSeriesHalfNulls(b *testing.B) {
	benchmarkInvert(b, 10, test.RandFloats10k, test.RandFloatsWithNulls10k)
}
func BenchmarkInvert10k_100SomeSeriesHalfNulls(b *testing.B) {
	benchmarkInvert(b, 100, test.RandFloats10k, test.RandFloatsWithNulls10k)
}
func BenchmarkInvert10k_1000SomeSeriesHalfNulls(b *testing.B) {
	benchmarkInvert(b, 1000, test.RandFloats10k, test.RandFloatsWithNulls10k)
}
func BenchmarkInvert10k_1AllSeriesHalfNulls(b *testing.B) {
	benchmarkInvert(b, 1, test.RandFloatsWithNulls10k, test.RandFloatsWithNulls10k)
}
func BenchmarkInvert10k_10AllSeriesHalfNulls(b *testing.B) {
	benchmarkInvert(b, 10, test.RandFloatsWithNulls10k, test.RandFloatsWithNulls10k)
}
func BenchmarkInvert10k_100AllSeriesHalfNulls(b *testing.B) {
	benchmarkInvert(b, 100, test.RandFloatsWithNulls10k, test.RandFloatsWithNulls10k)
}
func BenchmarkInvert10k_1000AllSeriesHalfNulls(b *testing.B) {
	benchmarkInvert(b, 1000, test.RandFloatsWithNulls10k, test.RandFloatsWithNulls10k)
}
func benchmarkInvert(b *testing.B, numSeries int, fn0, fn1 func() []schema.Point) {
	var input []models.Series
	for i := 0; i < numSeries; i++ {
		series := models.Series{
			QueryPatt: strconv.Itoa(i),
		}
		if i%2 == 0 {
			series.Datapoints = fn0()
		} else {
			series.Datapoints = fn1()
		}
		input = append(input, series)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		f := NewInvert()
		f.(*FuncInvert).in = NewMock(input)
		got, err := f.Exec(make(map[Req][]models.Series))
		if err != nil {
			b.Fatalf("%s", err)
		}
		results = got
	}
}
<commit_msg>make tests more consistent with new timeshift tests<commit_after>package expr

import (
	"math"
	"strconv"
	"testing"

	"github.com/grafana/metrictank/api/models"
	"github.com/grafana/metrictank/schema"
	"github.com/grafana/metrictank/test"
)

var datapoints = []schema.Point{
	{Val: 0, Ts: 10},
	{Val: math.NaN(), Ts: 20},
	{Val: -10, Ts: 30},
	{Val: 5.5, Ts: 40},
	{Val: -math.MaxFloat64, Ts: 50},
	{Val: -1234567890, Ts: 60},
	{Val: math.MaxFloat64, Ts: 70},
}

var datapointsInvert = []schema.Point{
	{Val: math.NaN(), Ts: 10},
	{Val: math.NaN(), Ts: 20},
	{Val: math.Pow(-10, -1), Ts: 30},
	{Val: math.Pow(5.5, -1), Ts: 40},
	{Val: math.Pow(-math.MaxFloat64, -1), Ts: 50},
	{Val: math.Pow(-1234567890, -1), Ts: 60},
	{Val: math.Pow(math.MaxFloat64, -1), Ts: 70},
}

func TestInvertBasic(t *testing.T) {
	in := []models.Series{
		{
			Interval:   10,
			QueryPatt:  "queryPattHere",
			Target:     "targetHere",
			Datapoints: getCopy(datapoints),
		},
	}

	f := getNewInvert(in)
	out := []models.Series{
		{
			Interval:   10,
			QueryPatt:  "invert(queryPattHere)",
			Target:     "invert(targetHere)",
			Datapoints: getCopy(datapointsInvert),
		},
	}

	got, err := f.Exec(make(map[Req][]models.Series))
	if err := equalOutput(out, got, nil, err); err != nil {
		t.Fatal(err)
	}
}

func TestInvertInputUnchanged(t *testing.T) {
	in := []models.Series{
		{
			Interval:   10,
			QueryPatt:  "queryPattHere",
			Target:     "targetHere",
			Datapoints: getCopy(datapoints),
		},
	}

	// store copy of the original input
	inCopy := make([]models.Series, len(in))
	copy(inCopy, in)

	f := getNewInvert(in)
	_, err := f.Exec(make(map[Req][]models.Series))

	// make sure input hasn't changed after call to invert
	if err := equalOutput(in, inCopy, nil, err); err != nil {
		t.Fatal(err)
	}
}

func TestInvertZero(t *testing.T) {
	in := []models.Series{
		{
			Interval:   10,
			QueryPatt:  "queryPattHere",
			Target:     "targetHere",
			Datapoints: []schema.Point{},
		},
	}

	f := getNewInvert(in)

	out := []models.Series{
		{
			Interval:   10,
			QueryPatt:  "invert(queryPattHere)",
			Target:     "invert(targetHere)",
			Datapoints: []schema.Point{},
		},
	}

	got, err := f.Exec(make(map[Req][]models.Series))
	if err := equalOutput(out, got, nil, err); err != nil {
		t.Fatal(err)
	}
}

func TestInvertMultiple(t *testing.T) {
	in := []models.Series{
		{
			Interval:   10,
			QueryPatt:  "queryPattHere",
			Target:     "targetHere",
			Datapoints: getCopy(datapoints),
		},
		{
			Interval:   20,
			QueryPatt:  "queryPattHere2",
			Target:     "targetHere2",
			Datapoints: getCopy(datapoints),
		},
	}

	f := getNewInvert(in)
	out := []models.Series{
		{
			Interval:   10,
			QueryPatt:  "invert(queryPattHere)",
			Target:     "invert(targetHere)",
			Datapoints: getCopy(datapointsInvert),
		},
		{
			Interval:   20,
			QueryPatt:  "invert(queryPattHere2)",
			Target:     "invert(targetHere2)",
			Datapoints: getCopy(datapointsInvert),
		},
	}

	got, err := f.Exec(make(map[Req][]models.Series))
	if err := equalOutput(out, got, nil, err); err != nil {
		t.Fatal(err)
	}
}

func getNewInvert(in []models.Series) *FuncInvert {
	f := NewInvert()
	ps := f.(*FuncInvert)
	ps.in = NewMock(in)
	return ps
}

func BenchmarkInvert10k_1NoNulls(b *testing.B) {
	benchmarkInvert(b, 1, test.RandFloats10k, test.RandFloats10k)
}
func BenchmarkInvert10k_10NoNulls(b *testing.B) {
	benchmarkInvert(b, 10, test.RandFloats10k, test.RandFloats10k)
}
func BenchmarkInvert10k_100NoNulls(b *testing.B) {
	benchmarkInvert(b, 100, test.RandFloats10k, test.RandFloats10k)
}
func BenchmarkInvert10k_1000NoNulls(b *testing.B) {
	benchmarkInvert(b, 1000, test.RandFloats10k, test.RandFloats10k)
}
func BenchmarkInvert10k_1SomeSeriesHalfNulls(b *testing.B) {
	benchmarkInvert(b, 1, test.RandFloats10k, test.RandFloatsWithNulls10k)
}
func BenchmarkInvert10k_10SomeSeriesHalfNulls(b *testing.B) {
	benchmarkInvert(b, 10, test.RandFloats10k, test.RandFloatsWithNulls10k)
}
func BenchmarkInvert10k_100SomeSeriesHalfNulls(b *testing.B) {
	benchmarkInvert(b, 100, test.RandFloats10k, test.RandFloatsWithNulls10k)
}
func BenchmarkInvert10k_1000SomeSeriesHalfNulls(b *testing.B) {
	benchmarkInvert(b, 1000, test.RandFloats10k, test.RandFloatsWithNulls10k)
}
func BenchmarkInvert10k_1AllSeriesHalfNulls(b *testing.B) {
	benchmarkInvert(b, 1, test.RandFloatsWithNulls10k, test.RandFloatsWithNulls10k)
}
func BenchmarkInvert10k_10AllSeriesHalfNulls(b *testing.B) {
	benchmarkInvert(b, 10, test.RandFloatsWithNulls10k, test.RandFloatsWithNulls10k)
}
func BenchmarkInvert10k_100AllSeriesHalfNulls(b *testing.B) {
	benchmarkInvert(b, 100, test.RandFloatsWithNulls10k, test.RandFloatsWithNulls10k)
}
func BenchmarkInvert10k_1000AllSeriesHalfNulls(b *testing.B) {
	benchmarkInvert(b, 1000, test.RandFloatsWithNulls10k, test.RandFloatsWithNulls10k)
}
func benchmarkInvert(b *testing.B, numSeries int, fn0, fn1 func() []schema.Point) {
	var input []models.Series
	for i := 0; i < numSeries; i++ {
		series := models.Series{
			QueryPatt: strconv.Itoa(i),
		}
		if i%2 == 0 {
			series.Datapoints = fn0()
		} else {
			series.Datapoints = fn1()
		}
		input = append(input, series)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		f := NewInvert()
		f.(*FuncInvert).in = NewMock(input)
		got, err := f.Exec(make(map[Req][]models.Series))
		if err != nil {
			b.Fatalf("%s", err)
		}
		results = got
	}
}
<|endoftext|>


<commit_before>package grafana

import (
	"encoding/json"
	"time"
)

type DashboardID uint64

type dashboardStyle string

const (
	dashboardDarkStyle  dashboardStyle = "dark"
	dashboardLightStyle dashboardStyle = "light"
)

type Dashboard struct {
	ID      DashboardID `json:"id"`
	Version uint64      `json:"version"`

	Editable      bool           `json:"editable"`
	GraphTooltip  uint8          `json:"graphTooltip"`
	HideControls  bool           `json:"hideControls"`
	Rows          []*Row         `json:"rows"`
	Style         dashboardStyle `json:"style"`
	Timezone      string         `json:"timezone"`
	Title         string         `json:"title"`
	tags          []string
	Meta          *DashboardMeta `json:"meta,omitempty"`
	SchemaVersion int            `json:"schemaVersion"`
}

// NewDashboard creates new Dashboard.
func NewDashboard(title string) *Dashboard {
	return &Dashboard{
		Title:         title,
		Editable:      true,
		SchemaVersion: 14,
		Style:         dashboardDarkStyle,
	}
}

func (d *Dashboard) String() string {
	return Stringify(d)
}

// Tags is a getter for Dashboard tags field
func (d *Dashboard) Tags() []string {
	return d.tags
}

// SetTags sets new tags to dashboard
func (d *Dashboard) SetTags(tags ...string) {
	newTags := []string{}
	uniqTags := make(map[string]bool)
	for _, tag := range tags {
		if _, ok := uniqTags[tag]; ok {
			continue
		}

		uniqTags[tag] = true
		newTags = append(newTags, tag)
	}

	d.tags = newTags
}

// AddTags adds given tags to dashboard. This method keeps uniqueness of tags.
func (d *Dashboard) AddTags(tags ...string) {
	tagFound := make(map[string]bool, len(d.tags))
	for _, tag := range d.tags {
		tagFound[tag] = true
	}

	for _, tag := range tags {
		if _, ok := tagFound[tag]; ok {
			continue
		}
		d.tags = append(d.tags, tag)
	}
}

// RemoveTags removes given tags from dashboard. Does nothing if tag is not found.
func (d *Dashboard) RemoveTags(tags ...string) {
	tagIndex := make(map[string]int, len(d.tags))
	for i, tag := range d.tags {
		tagIndex[tag] = i
	}

	for _, tag := range tags {
		if i, ok := tagIndex[tag]; ok {
			d.tags = append(d.tags[:i], d.tags[i+1:]...)
		}
	}
}

// UnmarshalJSON implements json.Unmarshaler interface
func (d *Dashboard) UnmarshalJSON(data []byte) error {
	type JSONDashboard Dashboard
	inDashboard := struct {
		*JSONDashboard
		Tags []string `json:"tags"`
	}{
		JSONDashboard: (*JSONDashboard)(d),
	}
	if err := json.Unmarshal(data, &inDashboard); err != nil {
		return err
	}

	d.tags = inDashboard.Tags

	return nil
}

// MarshalJSON implements json.Marshaler interface
func (d *Dashboard) MarshalJSON() ([]byte, error) {
	type JSONDashboard Dashboard
	dd := (*JSONDashboard)(d)
	dd.Meta = nil

	return json.Marshal(&struct {
		*JSONDashboard
		Tags []string `json:"tags,omitempty"`
	}{
		JSONDashboard: dd,
		Tags:          d.Tags(),
	})
}

type DashboardMeta struct {
	Slug    string `json:"slug"`
	Type    string `json:"type"`
	Version int    `json:"version"`

	CanEdit bool `json:"canEdit"`
	CanSave bool `json:"canSave"`
	CanStar bool `json:"canStar"`

	Created   time.Time `json:"created"`
	CreatedBy string    `json:"createdBy"`
	Expires   time.Time `json:"expires"`
	Updated   time.Time `json:"updated"`
	UpdatedBy string    `json:"updatedBy"`
}

func (dm *DashboardMeta) String() string {
	return Stringify(dm)
}

type Row struct {
	Collapsed bool         `json:"collapse"`
	Editable  bool         `json:"editable"`
	Height    string       `json:"height"`
	Panels    []*TextPanel `json:"panels"`
	RepeatFor string       `json:"repeat"` // repeat row for given variable
	ShowTitle bool         `json:"showTitle"`
	Title     string       `json:"title"`
	TitleSize string       `json:"titleSize"` // TODO: validation: h1-h6
}

// NewRow creates new Row with somw defaults.
func NewRow() *Row {
	return &Row{
		Editable: true,
	}
}

type textPanelMode string

const (
	TextPanelHTMLMode     textPanelMode = "html"
	TextPanelMarkdownMode textPanelMode = "markdown"
	TextPanelTextMode     textPanelMode = "text"
	textPanelDefaultMode  textPanelMode = TextPanelTextMode
)

type TextPanel struct {
	Content string        `json:"content"`
	Mode    textPanelMode `json:"mode"` // markdown/html/text, required

	// General options
	ID          int    `json:"id"` // Not sure if it's necessary
	Description string `json:"description"`
	Height      string `json:"height"`
	// Links       []*PanelLink `json:"links"`
	MinSpan     int    `json:"minSpan"`        // TODO: valid values: 1-12
	Span        int    `json:"span,omitempty"` // TODO: valid values: 1-12
	Title       string `json:"title"`
	Transparent bool   `json:"transparent"`
	Type        string `json:"type"` // required
}

// NewTextPanel creates new "Text" panel.
func NewTextPanel() *TextPanel {
	return &TextPanel{
		Mode: textPanelDefaultMode,
	}
}

type PanelLink struct {
	IncludeVars  bool   `json:"includeVars"`
	KeepTime     bool   `json:"keepTime"`
	Params       string `json:"params"`
	OpenInNewTab bool   `json:"targetBlank"`
	Type         string `json:"type"` // TODO validation: absolute/dashboard

	// type=absolute
	Title string `json:"title,omitempty"`
	URL   string `json:"url,omitempty"`

	// type=dashboard
	DashboardURI string `json:"dashUri,omitempty"`   // TODO: validation. should be valid dashboard
	Dashboard    string `json:"dashboard,omitempty"` // actually it's title
}

// NewPanelLink creates new PanelLink
func NewPanelLink(panelType string) *PanelLink {
	return &PanelLink{
		Type: panelType, // TODO: validation
	}
}
<commit_msg>Don't json marshall Meta field of Dashboard<commit_after>package grafana

import (
	"encoding/json"
	"time"
)

type DashboardID uint64

type dashboardStyle string

const (
	dashboardDarkStyle  dashboardStyle = "dark"
	dashboardLightStyle dashboardStyle = "light"
)

type Dashboard struct {
	ID      DashboardID `json:"id"`
	Version uint64      `json:"version"`

	Editable      bool           `json:"editable"`
	GraphTooltip  uint8          `json:"graphTooltip"`
	HideControls  bool           `json:"hideControls"`
	Rows          []*Row         `json:"rows"`
	Style         dashboardStyle `json:"style"`
	Timezone      string         `json:"timezone"`
	Title         string         `json:"title"`
	tags          []string
	Meta          *DashboardMeta `json:"meta,omitempty"`
	SchemaVersion int            `json:"schemaVersion"`
}

// NewDashboard creates new Dashboard.
func NewDashboard(title string) *Dashboard {
	return &Dashboard{
		Title:         title,
		Editable:      true,
		SchemaVersion: 14,
		Style:         dashboardDarkStyle,
	}
}

func (d *Dashboard) String() string {
	return Stringify(d)
}

// Tags is a getter for Dashboard tags field
func (d *Dashboard) Tags() []string {
	return d.tags
}

// SetTags sets new tags to dashboard
func (d *Dashboard) SetTags(tags ...string) {
	newTags := []string{}
	uniqTags := make(map[string]bool)
	for _, tag := range tags {
		if _, ok := uniqTags[tag]; ok {
			continue
		}

		uniqTags[tag] = true
		newTags = append(newTags, tag)
	}

	d.tags = newTags
}

// AddTags adds given tags to dashboard. This method keeps uniqueness of tags.
func (d *Dashboard) AddTags(tags ...string) {
	tagFound := make(map[string]bool, len(d.tags))
	for _, tag := range d.tags {
		tagFound[tag] = true
	}

	for _, tag := range tags {
		if _, ok := tagFound[tag]; ok {
			continue
		}
		d.tags = append(d.tags, tag)
	}
}

// RemoveTags removes given tags from dashboard. Does nothing if tag is not found.
func (d *Dashboard) RemoveTags(tags ...string) {
	tagIndex := make(map[string]int, len(d.tags))
	for i, tag := range d.tags {
		tagIndex[tag] = i
	}

	for _, tag := range tags {
		if i, ok := tagIndex[tag]; ok {
			d.tags = append(d.tags[:i], d.tags[i+1:]...)
		}
	}
}

// UnmarshalJSON implements json.Unmarshaler interface
func (d *Dashboard) UnmarshalJSON(data []byte) error {
	type JSONDashboard Dashboard
	inDashboard := struct {
		*JSONDashboard
		Tags []string `json:"tags"`
	}{
		JSONDashboard: (*JSONDashboard)(d),
	}
	if err := json.Unmarshal(data, &inDashboard); err != nil {
		return err
	}

	d.tags = inDashboard.Tags

	return nil
}

// MarshalJSON implements json.Marshaler interface
func (d *Dashboard) MarshalJSON() ([]byte, error) {
	type JSONDashboard Dashboard
	dd := (*JSONDashboard)(d)
	dd.Meta = nil

	return json.Marshal(&struct {
		*JSONDashboard
		Tags []string       `json:"tags"`
		Meta *DashboardMeta `json:"-"`
	}{
		JSONDashboard: dd,
		Tags:          d.Tags(),
	})
}

type DashboardMeta struct {
	Slug    string `json:"slug"`
	Type    string `json:"type"`
	Version int    `json:"version"`

	CanEdit bool `json:"canEdit"`
	CanSave bool `json:"canSave"`
	CanStar bool `json:"canStar"`

	Created   time.Time `json:"created"`
	CreatedBy string    `json:"createdBy"`
	Expires   time.Time `json:"expires"`
	Updated   time.Time `json:"updated"`
	UpdatedBy string    `json:"updatedBy"`
}

func (dm *DashboardMeta) String() string {
	return Stringify(dm)
}

type Row struct {
	Collapsed bool         `json:"collapse"`
	Editable  bool         `json:"editable"`
	Height    string       `json:"height"`
	Panels    []*TextPanel `json:"panels"`
	RepeatFor string       `json:"repeat"` // repeat row for given variable
	ShowTitle bool         `json:"showTitle"`
	Title     string       `json:"title"`
	TitleSize string       `json:"titleSize"` // TODO: validation: h1-h6
}

// NewRow creates new Row with somw defaults.
func NewRow() *Row {
	return &Row{
		Editable: true,
	}
}

type textPanelMode string

const (
	TextPanelHTMLMode     textPanelMode = "html"
	TextPanelMarkdownMode textPanelMode = "markdown"
	TextPanelTextMode     textPanelMode = "text"
	textPanelDefaultMode  textPanelMode = TextPanelTextMode
)

type TextPanel struct {
	Content string        `json:"content"`
	Mode    textPanelMode `json:"mode"` // markdown/html/text, required

	// General options
	ID          int    `json:"id"` // Not sure if it's necessary
	Description string `json:"description"`
	Height      string `json:"height"`
	// Links       []*PanelLink `json:"links"`
	MinSpan     int    `json:"minSpan"`        // TODO: valid values: 1-12
	Span        int    `json:"span,omitempty"` // TODO: valid values: 1-12
	Title       string `json:"title"`
	Transparent bool   `json:"transparent"`
	Type        string `json:"type"` // required
}

// NewTextPanel creates new "Text" panel.
func NewTextPanel() *TextPanel {
	return &TextPanel{
		Mode: textPanelDefaultMode,
	}
}

type PanelLink struct {
	IncludeVars  bool   `json:"includeVars"`
	KeepTime     bool   `json:"keepTime"`
	Params       string `json:"params"`
	OpenInNewTab bool   `json:"targetBlank"`
	Type         string `json:"type"` // TODO validation: absolute/dashboard

	// type=absolute
	Title string `json:"title,omitempty"`
	URL   string `json:"url,omitempty"`

	// type=dashboard
	DashboardURI string `json:"dashUri,omitempty"`   // TODO: validation. should be valid dashboard
	Dashboard    string `json:"dashboard,omitempty"` // actually it's title
}

// NewPanelLink creates new PanelLink
func NewPanelLink(panelType string) *PanelLink {
	return &PanelLink{
		Type: panelType, // TODO: validation
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2013 CoreOS Inc.

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 dbus

import (
	"errors"
	"github.com/godbus/dbus"
)

func (c *Conn) initJobs() {
	c.jobListener.jobs = make(map[dbus.ObjectPath]chan string)
}

func (c *Conn) jobComplete(signal *dbus.Signal) {
	var id uint32
	var job dbus.ObjectPath
	var unit string
	var result string
	dbus.Store(signal.Body, &id, &job, &unit, &result)
	c.jobListener.Lock()
	out, ok := c.jobListener.jobs[job]
	if ok {
		out <- result
		delete(c.jobListener.jobs, job)
	}
	c.jobListener.Unlock()
}

func (c *Conn) startJob(job string, args ...interface{}) (<-chan string, error) {
	c.jobListener.Lock()
	defer c.jobListener.Unlock()

	ch := make(chan string, 1)
	var path dbus.ObjectPath
	err := c.sysobj.Call(job, 0, args...).Store(&path)
	if err != nil {
		return nil, err
	}
	c.jobListener.jobs[path] = ch
	return ch, nil
}

func (c *Conn) runJob(job string, args ...interface{}) (string, error) {
	respCh, err := c.startJob(job, args...)
	if err != nil {
		return "", err
	}
	return <-respCh, nil
}

// StartUnit enqeues a start job and depending jobs, if any (unless otherwise
// specified by the mode string).
//
// Takes the unit to activate, plus a mode string. The mode needs to be one of
// replace, fail, isolate, ignore-dependencies, ignore-requirements. If
// "replace" the call will start the unit and its dependencies, possibly
// replacing already queued jobs that conflict with this. If "fail" the call
// will start the unit and its dependencies, but will fail if this would change
// an already queued job. If "isolate" the call will start the unit in question
// and terminate all units that aren't dependencies of it. If
// "ignore-dependencies" it will start a unit but ignore all its dependencies.
// If "ignore-requirements" it will start a unit but only ignore the
// requirement dependencies. It is not recommended to make use of the latter
// two options.
//
// Result string: one of done, canceled, timeout, failed, dependency, skipped.
// done indicates successful execution of a job. canceled indicates that a job
// has been canceled  before it finished execution. timeout indicates that the
// job timeout was reached. failed indicates that the job failed. dependency
// indicates that a job this job has been depending on failed and the job hence
// has been removed too. skipped indicates that a job was skipped because it
// didn't apply to the units current state.
func (c *Conn) StartUnit(name string, mode string) (string, error) {
	return c.runJob("org.freedesktop.systemd1.Manager.StartUnit", name, mode)
}

// StopUnit is similar to StartUnit but stops the specified unit rather
// than starting it.
func (c *Conn) StopUnit(name string, mode string) (string, error) {
	return c.runJob("org.freedesktop.systemd1.Manager.StopUnit", name, mode)
}

// ReloadUnit reloads a unit.  Reloading is done only if the unit is already running and fails otherwise.
func (c *Conn) ReloadUnit(name string, mode string) (string, error) {
	return c.runJob("org.freedesktop.systemd1.Manager.ReloadUnit", name, mode)
}

// RestartUnit restarts a service.  If a service is restarted that isn't
// running it will be started.
func (c *Conn) RestartUnit(name string, mode string) (string, error) {
	return c.runJob("org.freedesktop.systemd1.Manager.RestartUnit", name, mode)
}

// TryRestartUnit is like RestartUnit, except that a service that isn't running
// is not affected by the restart.
func (c *Conn) TryRestartUnit(name string, mode string) (string, error) {
	return c.runJob("org.freedesktop.systemd1.Manager.TryRestartUnit", name, mode)
}

// ReloadOrRestart attempts a reload if the unit supports it and use a restart
// otherwise.
func (c *Conn) ReloadOrRestartUnit(name string, mode string) (string, error) {
	return c.runJob("org.freedesktop.systemd1.Manager.ReloadOrRestartUnit", name, mode)
}

// ReloadOrTryRestart attempts a reload if the unit supports it and use a "Try"
// flavored restart otherwise.
func (c *Conn) ReloadOrTryRestartUnit(name string, mode string) (string, error) {
	return c.runJob("org.freedesktop.systemd1.Manager.ReloadOrTryRestartUnit", name, mode)
}

// StartTransientUnit() may be used to create and start a transient unit, which
// will be released as soon as it is not running or referenced anymore or the
// system is rebooted. name is the unit name including suffix, and must be
// unique. mode is the same as in StartUnit(), properties contains properties
// of the unit.
func (c *Conn) StartTransientUnit(name string, mode string, properties ...Property) (string, error) {
	return c.runJob("org.freedesktop.systemd1.Manager.StartTransientUnit", name, mode, properties, make([]PropertyCollection, 0))
}

// KillUnit takes the unit name and a UNIX signal number to send.  All of the unit's
// processes are killed.
func (c *Conn) KillUnit(name string, signal int32) {
	c.sysobj.Call("org.freedesktop.systemd1.Manager.KillUnit", 0, name, "all", signal).Store()
}

// getProperties takes the unit name and returns all of its dbus object properties, for the given dbus interface
func (c *Conn) getProperties(unit string, dbusInterface string) (map[string]interface{}, error) {
	var err error
	var props map[string]dbus.Variant

	path := ObjectPath("/org/freedesktop/systemd1/unit/" + unit)
	if !path.IsValid() {
		return nil, errors.New("invalid unit name: " + unit)
	}

	obj := c.sysconn.Object("org.freedesktop.systemd1", path)
	err = obj.Call("org.freedesktop.DBus.Properties.GetAll", 0, dbusInterface).Store(&props)
	if err != nil {
		return nil, err
	}

	out := make(map[string]interface{}, len(props))
	for k, v := range props {
		out[k] = v.Value()
	}

	return out, nil
}

// GetUnitProperties takes the unit name and returns all of its dbus object properties.
func (c *Conn) GetUnitProperties(unit string) (map[string]interface{}, error) {
	return c.getProperties(unit, "org.freedesktop.systemd1.Unit")
}

func (c *Conn) getProperty(unit string, dbusInterface string, propertyName string) (*Property, error) {
	var err error
	var prop dbus.Variant

	path := ObjectPath("/org/freedesktop/systemd1/unit/" + unit)
	if !path.IsValid() {
		return nil, errors.New("invalid unit name: " + unit)
	}

	obj := c.sysconn.Object("org.freedesktop.systemd1", path)
	err = obj.Call("org.freedesktop.DBus.Properties.Get", 0, dbusInterface, propertyName).Store(&prop)
	if err != nil {
		return nil, err
	}

	return &Property{Name: propertyName, Value: prop}, nil
}

func (c *Conn) GetUnitProperty(unit string, propertyName string) (*Property, error) {
	return c.getProperty(unit, "org.freedesktop.systemd1.Unit", propertyName)
}

// GetUnitTypeProperties returns the extra properties for a unit, specific to the unit type.
// Valid values for unitType: Service, Socket, Target, Device, Mount, Automount, Snapshot, Timer, Swap, Path, Slice, Scope
// return "dbus.Error: Unknown interface" if the unitType is not the correct type of the unit
func (c *Conn) GetUnitTypeProperties(unit string, unitType string) (map[string]interface{}, error) {
	return c.getProperties(unit, "org.freedesktop.systemd1."+unitType)
}

// SetUnitProperties() may be used to modify certain unit properties at runtime.
// Not all properties may be changed at runtime, but many resource management
// settings (primarily those in systemd.cgroup(5)) may. The changes are applied
// instantly, and stored on disk for future boots, unless runtime is true, in which
// case the settings only apply until the next reboot. name is the name of the unit
// to modify. properties are the settings to set, encoded as an array of property
// name and value pairs.
func (c *Conn) SetUnitProperties(name string, runtime bool, properties ...Property) error {
	return c.sysobj.Call("SetUnitProperties", 0, name, runtime, properties).Store()
}

func (c *Conn) GetUnitTypeProperty(unit string, unitType string, propertyName string) (*Property, error) {
	return c.getProperty(unit, "org.freedesktop.systemd1." + unitType, propertyName)
}

// ListUnits returns an array with all currently loaded units. Note that
// units may be known by multiple names at the same time, and hence there might
// be more unit names loaded than actual units behind them.
func (c *Conn) ListUnits() ([]UnitStatus, error) {
	result := make([][]interface{}, 0)
	err := c.sysobj.Call("org.freedesktop.systemd1.Manager.ListUnits", 0).Store(&result)
	if err != nil {
		return nil, err
	}

	resultInterface := make([]interface{}, len(result))
	for i := range result {
		resultInterface[i] = result[i]
	}

	status := make([]UnitStatus, len(result))
	statusInterface := make([]interface{}, len(status))
	for i := range status {
		statusInterface[i] = &status[i]
	}

	err = dbus.Store(resultInterface, statusInterface...)
	if err != nil {
		return nil, err
	}

	return status, nil
}

type UnitStatus struct {
	Name        string          // The primary unit name as string
	Description string          // The human readable description string
	LoadState   string          // The load state (i.e. whether the unit file has been loaded successfully)
	ActiveState string          // The active state (i.e. whether the unit is currently started or not)
	SubState    string          // The sub state (a more fine-grained version of the active state that is specific to the unit type, which the active state is not)
	Followed    string          // A unit that is being followed in its state by this unit, if there is any, otherwise the empty string.
	Path        dbus.ObjectPath // The unit object path
	JobId       uint32          // If there is a job queued for the job unit the numeric job id, 0 otherwise
	JobType     string          // The job type as string
	JobPath     dbus.ObjectPath // The job object path
}

// EnableUnitFiles() may be used to enable one or more units in the system (by
// creating symlinks to them in /etc or /run).
//
// It takes a list of unit files to enable (either just file names or full
// absolute paths if the unit files are residing outside the usual unit
// search paths), and two booleans: the first controls whether the unit shall
// be enabled for runtime only (true, /run), or persistently (false, /etc).
// The second one controls whether symlinks pointing to other units shall
// be replaced if necessary.
//
// This call returns one boolean and an array with the changes made. The
// boolean signals whether the unit files contained any enablement
// information (i.e. an [Install]) section. The changes list consists of
// structures with three strings: the type of the change (one of symlink
// or unlink), the file name of the symlink and the destination of the
// symlink.
func (c *Conn) EnableUnitFiles(files []string, runtime bool, force bool) (bool, []EnableUnitFileChange, error) {
	var carries_install_info bool

	result := make([][]interface{}, 0)
	err := c.sysobj.Call("org.freedesktop.systemd1.Manager.EnableUnitFiles", 0, files, runtime, force).Store(&carries_install_info, &result)
	if err != nil {
		return false, nil, err
	}

	resultInterface := make([]interface{}, len(result))
	for i := range result {
		resultInterface[i] = result[i]
	}

	changes := make([]EnableUnitFileChange, len(result))
	changesInterface := make([]interface{}, len(changes))
	for i := range changes {
		changesInterface[i] = &changes[i]
	}

	err = dbus.Store(resultInterface, changesInterface...)
	if err != nil {
		return false, nil, err
	}

	return carries_install_info, changes, nil
}

type EnableUnitFileChange struct {
	Type        string // Type of the change (one of symlink or unlink)
	Filename    string // File name of the symlink
	Destination string // Destination of the symlink
}

// DisableUnitFiles() may be used to disable one or more units in the system (by
// removing symlinks to them from /etc or /run).
//
// It takes a list of unit files to disable (either just file names or full
// absolute paths if the unit files are residing outside the usual unit
// search paths), and one boolean: whether the unit was enabled for runtime
// only (true, /run), or persistently (false, /etc).
//
// This call returns an array with the changes made. The changes list
// consists of structures with three strings: the type of the change (one of
// symlink or unlink), the file name of the symlink and the destination of the
// symlink.
func (c *Conn) DisableUnitFiles(files []string, runtime bool) ([]DisableUnitFileChange, error) {
	result := make([][]interface{}, 0)
	err := c.sysobj.Call("DisableUnitFiles", 0, files, runtime).Store(&result)
	if err != nil {
		return nil, err
	}

	resultInterface := make([]interface{}, len(result))
	for i := range result {
		resultInterface[i] = result[i]
	}

	changes := make([]DisableUnitFileChange, len(result))
	changesInterface := make([]interface{}, len(changes))
	for i := range changes {
		changesInterface[i] = &changes[i]
	}

	err = dbus.Store(resultInterface, changesInterface...)
	if err != nil {
		return nil, err
	}

	return changes, nil
}

type DisableUnitFileChange struct {
	Type        string // Type of the change (one of symlink or unlink)
	Filename    string // File name of the symlink
	Destination string // Destination of the symlink
}

// Reload instructs systemd to scan for and reload unit files. This is
// equivalent to a 'systemctl daemon-reload'.
func (c *Conn) Reload() error {
	return c.sysobj.Call("org.freedesktop.systemd1.Manager.Reload", 0).Store()
}
<commit_msg>fix(dbus): Use full path for SetUnitProperties<commit_after>/*
Copyright 2013 CoreOS Inc.

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 dbus

import (
	"errors"
	"github.com/godbus/dbus"
)

func (c *Conn) initJobs() {
	c.jobListener.jobs = make(map[dbus.ObjectPath]chan string)
}

func (c *Conn) jobComplete(signal *dbus.Signal) {
	var id uint32
	var job dbus.ObjectPath
	var unit string
	var result string
	dbus.Store(signal.Body, &id, &job, &unit, &result)
	c.jobListener.Lock()
	out, ok := c.jobListener.jobs[job]
	if ok {
		out <- result
		delete(c.jobListener.jobs, job)
	}
	c.jobListener.Unlock()
}

func (c *Conn) startJob(job string, args ...interface{}) (<-chan string, error) {
	c.jobListener.Lock()
	defer c.jobListener.Unlock()

	ch := make(chan string, 1)
	var path dbus.ObjectPath
	err := c.sysobj.Call(job, 0, args...).Store(&path)
	if err != nil {
		return nil, err
	}
	c.jobListener.jobs[path] = ch
	return ch, nil
}

func (c *Conn) runJob(job string, args ...interface{}) (string, error) {
	respCh, err := c.startJob(job, args...)
	if err != nil {
		return "", err
	}
	return <-respCh, nil
}

// StartUnit enqeues a start job and depending jobs, if any (unless otherwise
// specified by the mode string).
//
// Takes the unit to activate, plus a mode string. The mode needs to be one of
// replace, fail, isolate, ignore-dependencies, ignore-requirements. If
// "replace" the call will start the unit and its dependencies, possibly
// replacing already queued jobs that conflict with this. If "fail" the call
// will start the unit and its dependencies, but will fail if this would change
// an already queued job. If "isolate" the call will start the unit in question
// and terminate all units that aren't dependencies of it. If
// "ignore-dependencies" it will start a unit but ignore all its dependencies.
// If "ignore-requirements" it will start a unit but only ignore the
// requirement dependencies. It is not recommended to make use of the latter
// two options.
//
// Result string: one of done, canceled, timeout, failed, dependency, skipped.
// done indicates successful execution of a job. canceled indicates that a job
// has been canceled  before it finished execution. timeout indicates that the
// job timeout was reached. failed indicates that the job failed. dependency
// indicates that a job this job has been depending on failed and the job hence
// has been removed too. skipped indicates that a job was skipped because it
// didn't apply to the units current state.
func (c *Conn) StartUnit(name string, mode string) (string, error) {
	return c.runJob("org.freedesktop.systemd1.Manager.StartUnit", name, mode)
}

// StopUnit is similar to StartUnit but stops the specified unit rather
// than starting it.
func (c *Conn) StopUnit(name string, mode string) (string, error) {
	return c.runJob("org.freedesktop.systemd1.Manager.StopUnit", name, mode)
}

// ReloadUnit reloads a unit.  Reloading is done only if the unit is already running and fails otherwise.
func (c *Conn) ReloadUnit(name string, mode string) (string, error) {
	return c.runJob("org.freedesktop.systemd1.Manager.ReloadUnit", name, mode)
}

// RestartUnit restarts a service.  If a service is restarted that isn't
// running it will be started.
func (c *Conn) RestartUnit(name string, mode string) (string, error) {
	return c.runJob("org.freedesktop.systemd1.Manager.RestartUnit", name, mode)
}

// TryRestartUnit is like RestartUnit, except that a service that isn't running
// is not affected by the restart.
func (c *Conn) TryRestartUnit(name string, mode string) (string, error) {
	return c.runJob("org.freedesktop.systemd1.Manager.TryRestartUnit", name, mode)
}

// ReloadOrRestart attempts a reload if the unit supports it and use a restart
// otherwise.
func (c *Conn) ReloadOrRestartUnit(name string, mode string) (string, error) {
	return c.runJob("org.freedesktop.systemd1.Manager.ReloadOrRestartUnit", name, mode)
}

// ReloadOrTryRestart attempts a reload if the unit supports it and use a "Try"
// flavored restart otherwise.
func (c *Conn) ReloadOrTryRestartUnit(name string, mode string) (string, error) {
	return c.runJob("org.freedesktop.systemd1.Manager.ReloadOrTryRestartUnit", name, mode)
}

// StartTransientUnit() may be used to create and start a transient unit, which
// will be released as soon as it is not running or referenced anymore or the
// system is rebooted. name is the unit name including suffix, and must be
// unique. mode is the same as in StartUnit(), properties contains properties
// of the unit.
func (c *Conn) StartTransientUnit(name string, mode string, properties ...Property) (string, error) {
	return c.runJob("org.freedesktop.systemd1.Manager.StartTransientUnit", name, mode, properties, make([]PropertyCollection, 0))
}

// KillUnit takes the unit name and a UNIX signal number to send.  All of the unit's
// processes are killed.
func (c *Conn) KillUnit(name string, signal int32) {
	c.sysobj.Call("org.freedesktop.systemd1.Manager.KillUnit", 0, name, "all", signal).Store()
}

// getProperties takes the unit name and returns all of its dbus object properties, for the given dbus interface
func (c *Conn) getProperties(unit string, dbusInterface string) (map[string]interface{}, error) {
	var err error
	var props map[string]dbus.Variant

	path := ObjectPath("/org/freedesktop/systemd1/unit/" + unit)
	if !path.IsValid() {
		return nil, errors.New("invalid unit name: " + unit)
	}

	obj := c.sysconn.Object("org.freedesktop.systemd1", path)
	err = obj.Call("org.freedesktop.DBus.Properties.GetAll", 0, dbusInterface).Store(&props)
	if err != nil {
		return nil, err
	}

	out := make(map[string]interface{}, len(props))
	for k, v := range props {
		out[k] = v.Value()
	}

	return out, nil
}

// GetUnitProperties takes the unit name and returns all of its dbus object properties.
func (c *Conn) GetUnitProperties(unit string) (map[string]interface{}, error) {
	return c.getProperties(unit, "org.freedesktop.systemd1.Unit")
}

func (c *Conn) getProperty(unit string, dbusInterface string, propertyName string) (*Property, error) {
	var err error
	var prop dbus.Variant

	path := ObjectPath("/org/freedesktop/systemd1/unit/" + unit)
	if !path.IsValid() {
		return nil, errors.New("invalid unit name: " + unit)
	}

	obj := c.sysconn.Object("org.freedesktop.systemd1", path)
	err = obj.Call("org.freedesktop.DBus.Properties.Get", 0, dbusInterface, propertyName).Store(&prop)
	if err != nil {
		return nil, err
	}

	return &Property{Name: propertyName, Value: prop}, nil
}

func (c *Conn) GetUnitProperty(unit string, propertyName string) (*Property, error) {
	return c.getProperty(unit, "org.freedesktop.systemd1.Unit", propertyName)
}

// GetUnitTypeProperties returns the extra properties for a unit, specific to the unit type.
// Valid values for unitType: Service, Socket, Target, Device, Mount, Automount, Snapshot, Timer, Swap, Path, Slice, Scope
// return "dbus.Error: Unknown interface" if the unitType is not the correct type of the unit
func (c *Conn) GetUnitTypeProperties(unit string, unitType string) (map[string]interface{}, error) {
	return c.getProperties(unit, "org.freedesktop.systemd1."+unitType)
}

// SetUnitProperties() may be used to modify certain unit properties at runtime.
// Not all properties may be changed at runtime, but many resource management
// settings (primarily those in systemd.cgroup(5)) may. The changes are applied
// instantly, and stored on disk for future boots, unless runtime is true, in which
// case the settings only apply until the next reboot. name is the name of the unit
// to modify. properties are the settings to set, encoded as an array of property
// name and value pairs.
func (c *Conn) SetUnitProperties(name string, runtime bool, properties ...Property) error {
	return c.sysobj.Call("org.freedesktop.systemd1.Manager.SetUnitProperties", 0, name, runtime, properties).Store()
}

func (c *Conn) GetUnitTypeProperty(unit string, unitType string, propertyName string) (*Property, error) {
	return c.getProperty(unit, "org.freedesktop.systemd1." + unitType, propertyName)
}

// ListUnits returns an array with all currently loaded units. Note that
// units may be known by multiple names at the same time, and hence there might
// be more unit names loaded than actual units behind them.
func (c *Conn) ListUnits() ([]UnitStatus, error) {
	result := make([][]interface{}, 0)
	err := c.sysobj.Call("org.freedesktop.systemd1.Manager.ListUnits", 0).Store(&result)
	if err != nil {
		return nil, err
	}

	resultInterface := make([]interface{}, len(result))
	for i := range result {
		resultInterface[i] = result[i]
	}

	status := make([]UnitStatus, len(result))
	statusInterface := make([]interface{}, len(status))
	for i := range status {
		statusInterface[i] = &status[i]
	}

	err = dbus.Store(resultInterface, statusInterface...)
	if err != nil {
		return nil, err
	}

	return status, nil
}

type UnitStatus struct {
	Name        string          // The primary unit name as string
	Description string          // The human readable description string
	LoadState   string          // The load state (i.e. whether the unit file has been loaded successfully)
	ActiveState string          // The active state (i.e. whether the unit is currently started or not)
	SubState    string          // The sub state (a more fine-grained version of the active state that is specific to the unit type, which the active state is not)
	Followed    string          // A unit that is being followed in its state by this unit, if there is any, otherwise the empty string.
	Path        dbus.ObjectPath // The unit object path
	JobId       uint32          // If there is a job queued for the job unit the numeric job id, 0 otherwise
	JobType     string          // The job type as string
	JobPath     dbus.ObjectPath // The job object path
}

// EnableUnitFiles() may be used to enable one or more units in the system (by
// creating symlinks to them in /etc or /run).
//
// It takes a list of unit files to enable (either just file names or full
// absolute paths if the unit files are residing outside the usual unit
// search paths), and two booleans: the first controls whether the unit shall
// be enabled for runtime only (true, /run), or persistently (false, /etc).
// The second one controls whether symlinks pointing to other units shall
// be replaced if necessary.
//
// This call returns one boolean and an array with the changes made. The
// boolean signals whether the unit files contained any enablement
// information (i.e. an [Install]) section. The changes list consists of
// structures with three strings: the type of the change (one of symlink
// or unlink), the file name of the symlink and the destination of the
// symlink.
func (c *Conn) EnableUnitFiles(files []string, runtime bool, force bool) (bool, []EnableUnitFileChange, error) {
	var carries_install_info bool

	result := make([][]interface{}, 0)
	err := c.sysobj.Call("org.freedesktop.systemd1.Manager.EnableUnitFiles", 0, files, runtime, force).Store(&carries_install_info, &result)
	if err != nil {
		return false, nil, err
	}

	resultInterface := make([]interface{}, len(result))
	for i := range result {
		resultInterface[i] = result[i]
	}

	changes := make([]EnableUnitFileChange, len(result))
	changesInterface := make([]interface{}, len(changes))
	for i := range changes {
		changesInterface[i] = &changes[i]
	}

	err = dbus.Store(resultInterface, changesInterface...)
	if err != nil {
		return false, nil, err
	}

	return carries_install_info, changes, nil
}

type EnableUnitFileChange struct {
	Type        string // Type of the change (one of symlink or unlink)
	Filename    string // File name of the symlink
	Destination string // Destination of the symlink
}

// DisableUnitFiles() may be used to disable one or more units in the system (by
// removing symlinks to them from /etc or /run).
//
// It takes a list of unit files to disable (either just file names or full
// absolute paths if the unit files are residing outside the usual unit
// search paths), and one boolean: whether the unit was enabled for runtime
// only (true, /run), or persistently (false, /etc).
//
// This call returns an array with the changes made. The changes list
// consists of structures with three strings: the type of the change (one of
// symlink or unlink), the file name of the symlink and the destination of the
// symlink.
func (c *Conn) DisableUnitFiles(files []string, runtime bool) ([]DisableUnitFileChange, error) {
	result := make([][]interface{}, 0)
	err := c.sysobj.Call("DisableUnitFiles", 0, files, runtime).Store(&result)
	if err != nil {
		return nil, err
	}

	resultInterface := make([]interface{}, len(result))
	for i := range result {
		resultInterface[i] = result[i]
	}

	changes := make([]DisableUnitFileChange, len(result))
	changesInterface := make([]interface{}, len(changes))
	for i := range changes {
		changesInterface[i] = &changes[i]
	}

	err = dbus.Store(resultInterface, changesInterface...)
	if err != nil {
		return nil, err
	}

	return changes, nil
}

type DisableUnitFileChange struct {
	Type        string // Type of the change (one of symlink or unlink)
	Filename    string // File name of the symlink
	Destination string // Destination of the symlink
}

// Reload instructs systemd to scan for and reload unit files. This is
// equivalent to a 'systemctl daemon-reload'.
func (c *Conn) Reload() error {
	return c.sysobj.Call("org.freedesktop.systemd1.Manager.Reload", 0).Store()
}
<|endoftext|>


<commit_before>// Copyright 2015-2016, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// 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 dbr

import (
	"fmt"

	"strings"

	"github.com/ugorji/go/codec"
)

var (
	_ codec.Selfer = (*NullString)(nil)
)

// GoString satisfies the interface fmt.GoStringer when using %#v in Printf methods.
// Returns
// 		dbr.NewNullString(`...`,bool)
func (ns NullString) GoString() string {
	if ns.Valid && strings.ContainsRune(ns.String, '`') {
		// `This is my`string`
		ns.String = strings.Join(strings.Split(ns.String, "`"), "`+\"`\"+`")
		// `This is my`+"`"+`string`
	}

	ns.String = "`" + ns.String + "`"
	if !ns.Valid {
		ns.String = "nil"
	}
	return fmt.Sprintf("dbr.NewNullString(%s)", ns.String)
}

// CodecEncodeSelf for ugorji.go codec package
func (n NullString) CodecEncodeSelf(e *codec.Encoder) {
	if err := e.Encode(n.String); err != nil {
		PkgLog.Debug("dbr.NullString.CodecEncodeSelf", "err", err, "n", n)
	}
}

// CodecDecodeSelf  for ugorji.go codec package @todo write test ... not sure if ok
func (n *NullString) CodecDecodeSelf(d *codec.Decoder) {
	if err := d.Decode(&n.String); err != nil {
		PkgLog.Debug("dbr.NullString.CodecDecodeSelf", "err", err, "n", n)
	}
	// think about empty string and Valid value ...
}

// CodecEncodeSelf for ugorji.go codec package
func (n *NullInt64) CodecEncodeSelf(e *codec.Encoder) {
	if err := e.Encode(n.Int64); err != nil {
		PkgLog.Debug("dbr.NullInt64.CodecEncodeSelf", "err", err, "n", n)
	}
}

// CodecDecodeSelf  for ugorji.go codec package @todo write test ... not sure if ok
func (n *NullInt64) CodecDecodeSelf(d *codec.Decoder) {
	if err := d.Decode(&n.Int64); err != nil {
		PkgLog.Debug("dbr.NullInt64.CodecDecodeSelf", "err", err, "n", n)
	}
}

// CodecEncodeSelf for ugorji.go codec package
func (n NullFloat64) CodecEncodeSelf(e *codec.Encoder) {
	if err := e.Encode(n.Float64); err != nil {
		PkgLog.Debug("dbr.NullFloat64.CodecEncodeSelf", "err", err, "n", n)
	}
}

// CodecDecodeSelf  for ugorji.go codec package @todo write test ... not sure if ok
func (n *NullFloat64) CodecDecodeSelf(d *codec.Decoder) {
	if err := d.Decode(&n.Float64); err != nil {
		PkgLog.Debug("dbr.NullFloat64.CodecDecodeSelf", "err", err, "n", n)
	}
}

// CodecEncodeSelf for ugorji.go codec package
func (n NullTime) CodecEncodeSelf(e *codec.Encoder) {
	if err := e.Encode(n.Time); err != nil {
		PkgLog.Debug("dbr.NullTime.CodecEncodeSelf", "err", err, "n", n)
	}
}

// CodecDecodeSelf  for ugorji.go codec package @todo write test ... not sure if ok
func (n *NullTime) CodecDecodeSelf(d *codec.Decoder) {
	if err := d.Decode(&n.Time); err != nil {
		PkgLog.Debug("dbr.NullTime.CodecDecodeSelf", "err", err, "n", n)
	}
}

// CodecEncodeSelf for ugorji.go codec package
func (n NullBool) CodecEncodeSelf(e *codec.Encoder) {
	if err := e.Encode(n.Bool); err != nil {
		PkgLog.Debug("dbr.NullBool.CodecEncodeSelf", "err", err, "n", n)
	}
}

// CodecDecodeSelf  for ugorji.go codec package @todo write test ... not sure if ok
func (n *NullBool) CodecDecodeSelf(d *codec.Decoder) {
	if err := d.Decode(&n.Bool); err != nil {
		PkgLog.Debug("dbr.NullBool.CodecDecodeSelf", "err", err, "n", n)
	}
}
<commit_msg>storage/dbr: Remove github.com/ugorji/go/codec because file size<commit_after>// Copyright 2015-2016, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// 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 dbr

import (
	"fmt"
	"strings"

	// "github.com/ugorji/go/codec"
)

// removed because the binary grows for 5MB. will be added once needed.
// var _ codec.Selfer = (*NullString)(nil)

// GoString satisfies the interface fmt.GoStringer when using %#v in Printf methods.
// Returns
// 		dbr.NewNullString(`...`,bool)
func (ns NullString) GoString() string {
	if ns.Valid && strings.ContainsRune(ns.String, '`') {
		// `This is my`string`
		ns.String = strings.Join(strings.Split(ns.String, "`"), "`+\"`\"+`")
		// `This is my`+"`"+`string`
	}

	ns.String = "`" + ns.String + "`"
	if !ns.Valid {
		ns.String = "nil"
	}
	return fmt.Sprintf("dbr.NewNullString(%s)", ns.String)
}

//// CodecEncodeSelf for ugorji.go codec package
//func (n NullString) CodecEncodeSelf(e *codec.Encoder) {
//	if err := e.Encode(n.String); err != nil {
//		PkgLog.Debug("dbr.NullString.CodecEncodeSelf", "err", err, "n", n)
//	}
//}
//
//// CodecDecodeSelf  for ugorji.go codec package @todo write test ... not sure if ok
//func (n *NullString) CodecDecodeSelf(d *codec.Decoder) {
//	if err := d.Decode(&n.String); err != nil {
//		PkgLog.Debug("dbr.NullString.CodecDecodeSelf", "err", err, "n", n)
//	}
//	// think about empty string and Valid value ...
//}
//
//// CodecEncodeSelf for ugorji.go codec package
//func (n *NullInt64) CodecEncodeSelf(e *codec.Encoder) {
//	if err := e.Encode(n.Int64); err != nil {
//		PkgLog.Debug("dbr.NullInt64.CodecEncodeSelf", "err", err, "n", n)
//	}
//}
//
//// CodecDecodeSelf  for ugorji.go codec package @todo write test ... not sure if ok
//func (n *NullInt64) CodecDecodeSelf(d *codec.Decoder) {
//	if err := d.Decode(&n.Int64); err != nil {
//		PkgLog.Debug("dbr.NullInt64.CodecDecodeSelf", "err", err, "n", n)
//	}
//}
//
//// CodecEncodeSelf for ugorji.go codec package
//func (n NullFloat64) CodecEncodeSelf(e *codec.Encoder) {
//	if err := e.Encode(n.Float64); err != nil {
//		PkgLog.Debug("dbr.NullFloat64.CodecEncodeSelf", "err", err, "n", n)
//	}
//}
//
//// CodecDecodeSelf  for ugorji.go codec package @todo write test ... not sure if ok
//func (n *NullFloat64) CodecDecodeSelf(d *codec.Decoder) {
//	if err := d.Decode(&n.Float64); err != nil {
//		PkgLog.Debug("dbr.NullFloat64.CodecDecodeSelf", "err", err, "n", n)
//	}
//}
//
//// CodecEncodeSelf for ugorji.go codec package
//func (n NullTime) CodecEncodeSelf(e *codec.Encoder) {
//	if err := e.Encode(n.Time); err != nil {
//		PkgLog.Debug("dbr.NullTime.CodecEncodeSelf", "err", err, "n", n)
//	}
//}
//
//// CodecDecodeSelf  for ugorji.go codec package @todo write test ... not sure if ok
//func (n *NullTime) CodecDecodeSelf(d *codec.Decoder) {
//	if err := d.Decode(&n.Time); err != nil {
//		PkgLog.Debug("dbr.NullTime.CodecDecodeSelf", "err", err, "n", n)
//	}
//}
//
//// CodecEncodeSelf for ugorji.go codec package
//func (n NullBool) CodecEncodeSelf(e *codec.Encoder) {
//	if err := e.Encode(n.Bool); err != nil {
//		PkgLog.Debug("dbr.NullBool.CodecEncodeSelf", "err", err, "n", n)
//	}
//}
//
//// CodecDecodeSelf  for ugorji.go codec package @todo write test ... not sure if ok
//func (n *NullBool) CodecDecodeSelf(d *codec.Decoder) {
//	if err := d.Decode(&n.Bool); err != nil {
//		PkgLog.Debug("dbr.NullBool.CodecDecodeSelf", "err", err, "n", n)
//	}
//}
<|endoftext|>


<commit_before>package local

import (
	"fmt"
	"sort"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/golang/glog"

	clientmodel "github.com/prometheus/client_golang/model"
)

const maxMappedFP = 1 << 30 // About 1B fingerprints reserved for mapping.

var separatorString = string([]byte{clientmodel.SeparatorByte})

// fpMappings maps original fingerprints to a map of string representations of
// metrics to the truly unique fingerprint.
type fpMappings map[clientmodel.Fingerprint]map[string]clientmodel.Fingerprint

// fpMapper is used to map fingerprints in order to work around fingerprint
// collisions.
type fpMapper struct {
	mtx      sync.RWMutex // Protects mappings.
	mappings fpMappings

	fpToSeries      *seriesMap
	p               *persistence
	highestMappedFP clientmodel.Fingerprint
}

// newFPMapper loads the collision map from the persistence and
// returns an fpMapper ready to use.
func newFPMapper(fpToSeries *seriesMap, p *persistence) (*fpMapper, error) {
	r := &fpMapper{
		fpToSeries: fpToSeries,
		p:          p,
	}
	mappings, nextFP, err := p.loadFPMappings()
	if err != nil {
		return nil, err
	}
	r.mappings = mappings
	r.highestMappedFP = nextFP
	return r, nil
}

// mapFP takes a raw fingerprint (as returned by Metrics.FastFingerprint) and
// returns a truly unique fingerprint. The caller must have locked the raw
// fingerprint.
//
// If an error is encountered, it is returned together with the unchanged raw
// fingerprint.
func (r *fpMapper) mapFP(fp clientmodel.Fingerprint, m clientmodel.Metric) (clientmodel.Fingerprint, error) {
	// First check if we are in the reserved FP space, in which case this is
	// automatically a collision that has to be mapped.
	if fp <= maxMappedFP {
		return r.maybeAddMapping(fp, m)
	}

	// Then check the most likely case: This fp belongs to a series that is
	// already in memory.
	s, ok := r.fpToSeries.get(fp)
	if ok {
		// FP exists in memory, but is it for the same metric?
		if m.Equal(s.metric) {
			// Yupp. We are done.
			return fp, nil
		}
		// Collision detected!
		return r.maybeAddMapping(fp, m)
	}
	// Metric is not in memory. Before doing the expensive archive lookup,
	// check if we have a mapping for this metric in place already.
	r.mtx.RLock()
	mappedFPs, fpAlreadyMapped := r.mappings[fp]
	r.mtx.RUnlock()
	if fpAlreadyMapped {
		// We indeed have mapped fp historically.
		ms := metricToUniqueString(m)
		// fp is locked by the caller, so no further locking of
		// 'collisions' required (it is specific to fp).
		mappedFP, ok := mappedFPs[ms]
		if ok {
			// Historical mapping found, return the mapped FP.
			return mappedFP, nil
		}
	}
	// If we are here, FP does not exist in memory and is either not mapped
	// at all, or existing mappings for FP are not for m. Check if we have
	// something for FP in the archive.
	archivedMetric, err := r.p.getArchivedMetric(fp)
	if err != nil {
		return fp, err
	}
	if archivedMetric != nil {
		// FP exists in archive, but is it for the same metric?
		if m.Equal(archivedMetric) {
			// Yupp. We are done.
			return fp, nil
		}
		// Collision detected!
		return r.maybeAddMapping(fp, m)
	}
	// As fp does not exist, neither in memory nor in archive, we can safely
	// keep it unmapped.
	return fp, nil
}

// maybeAddMapping is only used internally. It takes a detected collision and
// adds it to the collisions map if not yet there. In any case, it returns the
// truly unique fingerprint for the colliding metric.
func (r *fpMapper) maybeAddMapping(
	fp clientmodel.Fingerprint,
	collidingMetric clientmodel.Metric,
) (clientmodel.Fingerprint, error) {
	ms := metricToUniqueString(collidingMetric)
	r.mtx.RLock()
	mappedFPs, ok := r.mappings[fp]
	r.mtx.RUnlock()
	if ok {
		// fp is locked by the caller, so no further locking required.
		mappedFP, ok := mappedFPs[ms]
		if ok {
			return mappedFP, nil // Existing mapping.
		}
		// A new mapping has to be created.
		mappedFP = r.nextMappedFP()
		mappedFPs[ms] = mappedFP
		r.mtx.RLock()
		// Checkpoint mappings after each change.
		err := r.p.checkpointFPMappings(r.mappings)
		r.mtx.RUnlock()
		glog.Infof(
			"Collision detected for fingerprint %v, metric %v, mapping to new fingerprint %v.",
			fp, collidingMetric, mappedFP,
		)
		return mappedFP, err
	}
	// This is the first collision for fp.
	mappedFP := r.nextMappedFP()
	mappedFPs = map[string]clientmodel.Fingerprint{ms: mappedFP}
	r.mtx.Lock()
	r.mappings[fp] = mappedFPs
	// Checkpoint mappings after each change.
	err := r.p.checkpointFPMappings(r.mappings)
	r.mtx.Unlock()
	glog.Infof(
		"Collision detected for fingerprint %v, metric %v, mapping to new fingerprint %v.",
		fp, collidingMetric, mappedFP,
	)
	return mappedFP, err
}

func (r *fpMapper) nextMappedFP() clientmodel.Fingerprint {
	mappedFP := clientmodel.Fingerprint(atomic.AddUint64((*uint64)(&r.highestMappedFP), 1))
	if mappedFP > maxMappedFP {
		panic(fmt.Errorf("more than %v fingerprints mapped in collision detection", maxMappedFP))
	}
	return mappedFP
}

// metricToUniqueString turns a metric into a string in a reproducible and
// unique way, i.e. the same metric will always create the same string, and
// different metrics will always create different strings. In a way, it is the
// "ideal" fingerprint function, only that it is more expensive than the
// FastFingerprint function, and its result is not suitable as a key for maps
// and indexes as it might become really large, causing a lot of hashing effort
// in maps and a lot of storage overhead in indexes.
func metricToUniqueString(m clientmodel.Metric) string {
	parts := make([]string, 0, len(m))
	for ln, lv := range m {
		parts = append(parts, string(ln)+separatorString+string(lv))
	}
	sort.Strings(parts)
	return strings.Join(parts, separatorString)
}
<commit_msg>Reserve only ~1M FPs for the mapping.<commit_after>package local

import (
	"fmt"
	"sort"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/golang/glog"

	clientmodel "github.com/prometheus/client_golang/model"
)

const maxMappedFP = 1 << 20 // About 1M fingerprints reserved for mapping.

var separatorString = string([]byte{clientmodel.SeparatorByte})

// fpMappings maps original fingerprints to a map of string representations of
// metrics to the truly unique fingerprint.
type fpMappings map[clientmodel.Fingerprint]map[string]clientmodel.Fingerprint

// fpMapper is used to map fingerprints in order to work around fingerprint
// collisions.
type fpMapper struct {
	mtx      sync.RWMutex // Protects mappings.
	mappings fpMappings

	fpToSeries      *seriesMap
	p               *persistence
	highestMappedFP clientmodel.Fingerprint
}

// newFPMapper loads the collision map from the persistence and
// returns an fpMapper ready to use.
func newFPMapper(fpToSeries *seriesMap, p *persistence) (*fpMapper, error) {
	r := &fpMapper{
		fpToSeries: fpToSeries,
		p:          p,
	}
	mappings, nextFP, err := p.loadFPMappings()
	if err != nil {
		return nil, err
	}
	r.mappings = mappings
	r.highestMappedFP = nextFP
	return r, nil
}

// mapFP takes a raw fingerprint (as returned by Metrics.FastFingerprint) and
// returns a truly unique fingerprint. The caller must have locked the raw
// fingerprint.
//
// If an error is encountered, it is returned together with the unchanged raw
// fingerprint.
func (r *fpMapper) mapFP(fp clientmodel.Fingerprint, m clientmodel.Metric) (clientmodel.Fingerprint, error) {
	// First check if we are in the reserved FP space, in which case this is
	// automatically a collision that has to be mapped.
	if fp <= maxMappedFP {
		return r.maybeAddMapping(fp, m)
	}

	// Then check the most likely case: This fp belongs to a series that is
	// already in memory.
	s, ok := r.fpToSeries.get(fp)
	if ok {
		// FP exists in memory, but is it for the same metric?
		if m.Equal(s.metric) {
			// Yupp. We are done.
			return fp, nil
		}
		// Collision detected!
		return r.maybeAddMapping(fp, m)
	}
	// Metric is not in memory. Before doing the expensive archive lookup,
	// check if we have a mapping for this metric in place already.
	r.mtx.RLock()
	mappedFPs, fpAlreadyMapped := r.mappings[fp]
	r.mtx.RUnlock()
	if fpAlreadyMapped {
		// We indeed have mapped fp historically.
		ms := metricToUniqueString(m)
		// fp is locked by the caller, so no further locking of
		// 'collisions' required (it is specific to fp).
		mappedFP, ok := mappedFPs[ms]
		if ok {
			// Historical mapping found, return the mapped FP.
			return mappedFP, nil
		}
	}
	// If we are here, FP does not exist in memory and is either not mapped
	// at all, or existing mappings for FP are not for m. Check if we have
	// something for FP in the archive.
	archivedMetric, err := r.p.getArchivedMetric(fp)
	if err != nil {
		return fp, err
	}
	if archivedMetric != nil {
		// FP exists in archive, but is it for the same metric?
		if m.Equal(archivedMetric) {
			// Yupp. We are done.
			return fp, nil
		}
		// Collision detected!
		return r.maybeAddMapping(fp, m)
	}
	// As fp does not exist, neither in memory nor in archive, we can safely
	// keep it unmapped.
	return fp, nil
}

// maybeAddMapping is only used internally. It takes a detected collision and
// adds it to the collisions map if not yet there. In any case, it returns the
// truly unique fingerprint for the colliding metric.
func (r *fpMapper) maybeAddMapping(
	fp clientmodel.Fingerprint,
	collidingMetric clientmodel.Metric,
) (clientmodel.Fingerprint, error) {
	ms := metricToUniqueString(collidingMetric)
	r.mtx.RLock()
	mappedFPs, ok := r.mappings[fp]
	r.mtx.RUnlock()
	if ok {
		// fp is locked by the caller, so no further locking required.
		mappedFP, ok := mappedFPs[ms]
		if ok {
			return mappedFP, nil // Existing mapping.
		}
		// A new mapping has to be created.
		mappedFP = r.nextMappedFP()
		mappedFPs[ms] = mappedFP
		r.mtx.RLock()
		// Checkpoint mappings after each change.
		err := r.p.checkpointFPMappings(r.mappings)
		r.mtx.RUnlock()
		glog.Infof(
			"Collision detected for fingerprint %v, metric %v, mapping to new fingerprint %v.",
			fp, collidingMetric, mappedFP,
		)
		return mappedFP, err
	}
	// This is the first collision for fp.
	mappedFP := r.nextMappedFP()
	mappedFPs = map[string]clientmodel.Fingerprint{ms: mappedFP}
	r.mtx.Lock()
	r.mappings[fp] = mappedFPs
	// Checkpoint mappings after each change.
	err := r.p.checkpointFPMappings(r.mappings)
	r.mtx.Unlock()
	glog.Infof(
		"Collision detected for fingerprint %v, metric %v, mapping to new fingerprint %v.",
		fp, collidingMetric, mappedFP,
	)
	return mappedFP, err
}

func (r *fpMapper) nextMappedFP() clientmodel.Fingerprint {
	mappedFP := clientmodel.Fingerprint(atomic.AddUint64((*uint64)(&r.highestMappedFP), 1))
	if mappedFP > maxMappedFP {
		panic(fmt.Errorf("more than %v fingerprints mapped in collision detection", maxMappedFP))
	}
	return mappedFP
}

// metricToUniqueString turns a metric into a string in a reproducible and
// unique way, i.e. the same metric will always create the same string, and
// different metrics will always create different strings. In a way, it is the
// "ideal" fingerprint function, only that it is more expensive than the
// FastFingerprint function, and its result is not suitable as a key for maps
// and indexes as it might become really large, causing a lot of hashing effort
// in maps and a lot of storage overhead in indexes.
func metricToUniqueString(m clientmodel.Metric) string {
	parts := make([]string, 0, len(m))
	for ln, lv := range m {
		parts = append(parts, string(ln)+separatorString+string(lv))
	}
	sort.Strings(parts)
	return strings.Join(parts, separatorString)
}
<|endoftext|>


<commit_before>// Copyright 2014 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

package storage

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/google/btree"

	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/storage/engine"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/hlc"
	"github.com/cockroachdb/cockroach/util/leaktest"
	"github.com/cockroachdb/cockroach/util/log"
	"github.com/cockroachdb/cockroach/util/stop"
	"github.com/cockroachdb/cockroach/util/timeutil"
)

// Test implementation of a range set backed by btree.BTree.
type testRangeSet struct {
	sync.Mutex
	rangesByKey *btree.BTree
	visited     int
}

// newTestRangeSet creates a new range set that has the count number of ranges.
func newTestRangeSet(count int, t *testing.T) *testRangeSet {
	rs := &testRangeSet{rangesByKey: btree.New(64 /* degree */)}
	for i := 0; i < count; i++ {
		desc := &roachpb.RangeDescriptor{
			RangeID:  roachpb.RangeID(i),
			StartKey: roachpb.RKey(fmt.Sprintf("%03d", i)),
			EndKey:   roachpb.RKey(fmt.Sprintf("%03d", i+1)),
		}
		// Initialize the range stat so the scanner can use it.
		rng := &Replica{
			RangeID: desc.RangeID,
			stats: &rangeStats{
				rangeID: desc.RangeID,
				MVCCStats: engine.MVCCStats{
					KeyBytes:  1,
					ValBytes:  2,
					KeyCount:  1,
					LiveCount: 1,
				},
			},
		}
		if err := rng.setDesc(desc); err != nil {
			t.Fatal(err)
		}
		if exRngItem := rs.rangesByKey.ReplaceOrInsert(rng); exRngItem != nil {
			t.Fatalf("failed to insert range %s", rng)
		}
	}
	return rs
}

func (rs *testRangeSet) Visit(visitor func(*Replica) bool) {
	rs.Lock()
	defer rs.Unlock()
	rs.visited = 0
	rs.rangesByKey.Ascend(func(i btree.Item) bool {
		rs.visited++
		rs.Unlock()
		defer rs.Lock()
		return visitor(i.(*Replica))
	})
}

func (rs *testRangeSet) EstimatedCount() int {
	rs.Lock()
	defer rs.Unlock()
	count := rs.rangesByKey.Len() - rs.visited
	if count < 1 {
		count = 1
	}
	return count
}

// removeRange removes the i-th range from the range set.
func (rs *testRangeSet) remove(index int, t *testing.T) *Replica {
	endKey := roachpb.Key(fmt.Sprintf("%03d", index+1))
	rs.Lock()
	defer rs.Unlock()
	rng := rs.rangesByKey.Delete((rangeBTreeKey)(endKey))
	if rng == nil {
		t.Fatalf("failed to delete range of end key %s", endKey)
	}
	return rng.(*Replica)
}

// Test implementation of a range queue which adds range to an
// internal slice.
type testQueue struct {
	sync.Mutex // Protects ranges, done & processed count
	ranges     []*Replica
	done       bool
	processed  int
	disabled   bool
}

// setDisabled suspends processing of items from the queue.
func (tq *testQueue) setDisabled(d bool) {
	tq.Lock()
	defer tq.Unlock()
	tq.disabled = d
}

func (tq *testQueue) Start(clock *hlc.Clock, stopper *stop.Stopper) {
	stopper.RunWorker(func() {
		for {
			select {
			case <-time.After(1 * time.Millisecond):
				tq.Lock()
				if !tq.disabled && len(tq.ranges) > 0 {
					tq.ranges = tq.ranges[1:]
					tq.processed++
				}
				tq.Unlock()
			case <-stopper.ShouldStop():
				tq.Lock()
				tq.done = true
				tq.Unlock()
				return
			}
		}
	})
}

func (tq *testQueue) MaybeAdd(rng *Replica, now roachpb.Timestamp) {
	tq.Lock()
	defer tq.Unlock()
	if index := tq.indexOf(rng); index == -1 {
		tq.ranges = append(tq.ranges, rng)
	}
}

func (tq *testQueue) MaybeRemove(rng *Replica) {
	tq.Lock()
	defer tq.Unlock()
	if index := tq.indexOf(rng); index != -1 {
		tq.ranges = append(tq.ranges[:index], tq.ranges[index+1:]...)
	}
}

func (tq *testQueue) count() int {
	tq.Lock()
	defer tq.Unlock()
	return len(tq.ranges)
}

func (tq *testQueue) indexOf(rng *Replica) int {
	for i, r := range tq.ranges {
		if r == rng {
			return i
		}
	}
	return -1
}

func (tq *testQueue) isDone() bool {
	tq.Lock()
	defer tq.Unlock()
	return tq.done
}

// TestScannerAddToQueues verifies that ranges are added to and
// removed from multiple queues.
func TestScannerAddToQueues(t *testing.T) {
	defer leaktest.AfterTest(t)()
	const count = 3
	ranges := newTestRangeSet(count, t)
	q1, q2 := &testQueue{}, &testQueue{}
	// We don't want to actually consume entries from the queues during this test.
	q1.setDisabled(true)
	q2.setDisabled(true)
	s := newReplicaScanner(1*time.Millisecond, 0, ranges)
	s.AddQueues(q1, q2)
	mc := hlc.NewManualClock(0)
	clock := hlc.NewClock(mc.UnixNano)
	stopper := stop.NewStopper()

	// Start queue and verify that all ranges are added to both queues.
	s.Start(clock, stopper)
	util.SucceedsSoon(t, func() error {
		if q1.count() != count || q2.count() != count {
			return util.Errorf("q1 or q2 count != %d; got %d, %d", count, q1.count(), q2.count())
		}
		return nil
	})

	// Remove first range and verify it does not exist in either range.
	rng := ranges.remove(0, t)
	util.SucceedsSoon(t, func() error {
		// This is intentionally inside the loop, otherwise this test races as
		// our removal of the range may be processed before a stray re-queue.
		// Removing on each attempt makes sure we clean this up as we retry.
		s.RemoveReplica(rng)
		c1 := q1.count()
		c2 := q2.count()
		if c1 != count-1 || c2 != count-1 {
			return util.Errorf("q1 or q2 count != %d; got %d, %d", count-1, c1, c2)
		}
		return nil
	})

	// Stop scanner and verify both queues are stopped.
	stopper.Stop()
	if !q1.isDone() || !q2.isDone() {
		t.Errorf("expected all queues to stop; got %t, %t", q1.isDone(), q2.isDone())
	}
}

// TestScannerTiming verifies that ranges are scanned, regardless
// of how many, to match scanInterval.
func TestScannerTiming(t *testing.T) {
	defer leaktest.AfterTest(t)()
	const count = 3
	const runTime = 100 * time.Millisecond
	const maxError = 7500 * time.Microsecond
	durations := []time.Duration{
		15 * time.Millisecond,
		25 * time.Millisecond,
	}
	for i, duration := range durations {
		util.SucceedsSoon(t, func() error {
			ranges := newTestRangeSet(count, t)
			q := &testQueue{}
			s := newReplicaScanner(duration, 0, ranges)
			s.AddQueues(q)
			mc := hlc.NewManualClock(0)
			clock := hlc.NewClock(mc.UnixNano)
			stopper := stop.NewStopper()
			s.Start(clock, stopper)
			time.Sleep(runTime)
			stopper.Stop()

			avg := s.avgScan()
			log.Infof("%d: average scan: %s", i, avg)
			if avg.Nanoseconds()-duration.Nanoseconds() > maxError.Nanoseconds() ||
				duration.Nanoseconds()-avg.Nanoseconds() > maxError.Nanoseconds() {
				return util.Errorf("expected %s, got %s: exceeds max error of %s", duration, avg, maxError)
			}
			return nil
		})
	}
}

// TestScannerPaceInterval tests that paceInterval returns the correct interval.
func TestScannerPaceInterval(t *testing.T) {
	defer leaktest.AfterTest(t)()
	const count = 3
	durations := []time.Duration{
		30 * time.Millisecond,
		60 * time.Millisecond,
		500 * time.Millisecond,
	}
	// function logs an error when the actual value is not close
	// to the expected value
	logErrorWhenNotCloseTo := func(expected, actual time.Duration) {
		delta := 1 * time.Millisecond
		if actual < expected-delta || actual > expected+delta {
			t.Errorf("Expected duration %s, got %s", expected, actual)
		}
	}
	for _, duration := range durations {
		startTime := timeutil.Now()
		ranges := newTestRangeSet(count, t)
		s := newReplicaScanner(duration, 0, ranges)
		interval := s.paceInterval(startTime, startTime)
		logErrorWhenNotCloseTo(duration/count, interval)
		// The range set is empty
		ranges = newTestRangeSet(0, t)
		s = newReplicaScanner(duration, 0, ranges)
		interval = s.paceInterval(startTime, startTime)
		logErrorWhenNotCloseTo(duration, interval)
		ranges = newTestRangeSet(count, t)
		s = newReplicaScanner(duration, 0, ranges)
		// Move the present to duration time into the future
		interval = s.paceInterval(startTime, startTime.Add(duration))
		logErrorWhenNotCloseTo(0, interval)
	}
}

// TestScannerEmptyRangeSet verifies that an empty range set doesn't busy loop.
func TestScannerEmptyRangeSet(t *testing.T) {
	defer leaktest.AfterTest(t)()
	ranges := newTestRangeSet(0, t)
	q := &testQueue{}
	s := newReplicaScanner(1*time.Millisecond, 0, ranges)
	s.AddQueues(q)
	mc := hlc.NewManualClock(0)
	clock := hlc.NewClock(mc.UnixNano)
	stopper := stop.NewStopper()
	defer stopper.Stop()
	s.Start(clock, stopper)
	time.Sleep(3 * time.Millisecond)
	if count := s.Count(); count > 3 {
		t.Errorf("expected three loops; got %d", count)
	}
}
<commit_msg>storage: deflake TestScannerEmptyRangeSet<commit_after>// Copyright 2014 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

package storage

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/google/btree"

	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/storage/engine"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/hlc"
	"github.com/cockroachdb/cockroach/util/leaktest"
	"github.com/cockroachdb/cockroach/util/log"
	"github.com/cockroachdb/cockroach/util/stop"
	"github.com/cockroachdb/cockroach/util/timeutil"
)

// Test implementation of a range set backed by btree.BTree.
type testRangeSet struct {
	sync.Mutex
	rangesByKey *btree.BTree
	visited     int
}

// newTestRangeSet creates a new range set that has the count number of ranges.
func newTestRangeSet(count int, t *testing.T) *testRangeSet {
	rs := &testRangeSet{rangesByKey: btree.New(64 /* degree */)}
	for i := 0; i < count; i++ {
		desc := &roachpb.RangeDescriptor{
			RangeID:  roachpb.RangeID(i),
			StartKey: roachpb.RKey(fmt.Sprintf("%03d", i)),
			EndKey:   roachpb.RKey(fmt.Sprintf("%03d", i+1)),
		}
		// Initialize the range stat so the scanner can use it.
		rng := &Replica{
			RangeID: desc.RangeID,
			stats: &rangeStats{
				rangeID: desc.RangeID,
				MVCCStats: engine.MVCCStats{
					KeyBytes:  1,
					ValBytes:  2,
					KeyCount:  1,
					LiveCount: 1,
				},
			},
		}
		if err := rng.setDesc(desc); err != nil {
			t.Fatal(err)
		}
		if exRngItem := rs.rangesByKey.ReplaceOrInsert(rng); exRngItem != nil {
			t.Fatalf("failed to insert range %s", rng)
		}
	}
	return rs
}

func (rs *testRangeSet) Visit(visitor func(*Replica) bool) {
	rs.Lock()
	defer rs.Unlock()
	rs.visited = 0
	rs.rangesByKey.Ascend(func(i btree.Item) bool {
		rs.visited++
		rs.Unlock()
		defer rs.Lock()
		return visitor(i.(*Replica))
	})
}

func (rs *testRangeSet) EstimatedCount() int {
	rs.Lock()
	defer rs.Unlock()
	count := rs.rangesByKey.Len() - rs.visited
	if count < 1 {
		count = 1
	}
	return count
}

// removeRange removes the i-th range from the range set.
func (rs *testRangeSet) remove(index int, t *testing.T) *Replica {
	endKey := roachpb.Key(fmt.Sprintf("%03d", index+1))
	rs.Lock()
	defer rs.Unlock()
	rng := rs.rangesByKey.Delete((rangeBTreeKey)(endKey))
	if rng == nil {
		t.Fatalf("failed to delete range of end key %s", endKey)
	}
	return rng.(*Replica)
}

// Test implementation of a range queue which adds range to an
// internal slice.
type testQueue struct {
	sync.Mutex // Protects ranges, done & processed count
	ranges     []*Replica
	done       bool
	processed  int
	disabled   bool
}

// setDisabled suspends processing of items from the queue.
func (tq *testQueue) setDisabled(d bool) {
	tq.Lock()
	defer tq.Unlock()
	tq.disabled = d
}

func (tq *testQueue) Start(clock *hlc.Clock, stopper *stop.Stopper) {
	stopper.RunWorker(func() {
		for {
			select {
			case <-time.After(1 * time.Millisecond):
				tq.Lock()
				if !tq.disabled && len(tq.ranges) > 0 {
					tq.ranges = tq.ranges[1:]
					tq.processed++
				}
				tq.Unlock()
			case <-stopper.ShouldStop():
				tq.Lock()
				tq.done = true
				tq.Unlock()
				return
			}
		}
	})
}

func (tq *testQueue) MaybeAdd(rng *Replica, now roachpb.Timestamp) {
	tq.Lock()
	defer tq.Unlock()
	if index := tq.indexOf(rng); index == -1 {
		tq.ranges = append(tq.ranges, rng)
	}
}

func (tq *testQueue) MaybeRemove(rng *Replica) {
	tq.Lock()
	defer tq.Unlock()
	if index := tq.indexOf(rng); index != -1 {
		tq.ranges = append(tq.ranges[:index], tq.ranges[index+1:]...)
	}
}

func (tq *testQueue) count() int {
	tq.Lock()
	defer tq.Unlock()
	return len(tq.ranges)
}

func (tq *testQueue) indexOf(rng *Replica) int {
	for i, r := range tq.ranges {
		if r == rng {
			return i
		}
	}
	return -1
}

func (tq *testQueue) isDone() bool {
	tq.Lock()
	defer tq.Unlock()
	return tq.done
}

// TestScannerAddToQueues verifies that ranges are added to and
// removed from multiple queues.
func TestScannerAddToQueues(t *testing.T) {
	defer leaktest.AfterTest(t)()
	const count = 3
	ranges := newTestRangeSet(count, t)
	q1, q2 := &testQueue{}, &testQueue{}
	// We don't want to actually consume entries from the queues during this test.
	q1.setDisabled(true)
	q2.setDisabled(true)
	s := newReplicaScanner(1*time.Millisecond, 0, ranges)
	s.AddQueues(q1, q2)
	mc := hlc.NewManualClock(0)
	clock := hlc.NewClock(mc.UnixNano)
	stopper := stop.NewStopper()

	// Start queue and verify that all ranges are added to both queues.
	s.Start(clock, stopper)
	util.SucceedsSoon(t, func() error {
		if q1.count() != count || q2.count() != count {
			return util.Errorf("q1 or q2 count != %d; got %d, %d", count, q1.count(), q2.count())
		}
		return nil
	})

	// Remove first range and verify it does not exist in either range.
	rng := ranges.remove(0, t)
	util.SucceedsSoon(t, func() error {
		// This is intentionally inside the loop, otherwise this test races as
		// our removal of the range may be processed before a stray re-queue.
		// Removing on each attempt makes sure we clean this up as we retry.
		s.RemoveReplica(rng)
		c1 := q1.count()
		c2 := q2.count()
		if c1 != count-1 || c2 != count-1 {
			return util.Errorf("q1 or q2 count != %d; got %d, %d", count-1, c1, c2)
		}
		return nil
	})

	// Stop scanner and verify both queues are stopped.
	stopper.Stop()
	if !q1.isDone() || !q2.isDone() {
		t.Errorf("expected all queues to stop; got %t, %t", q1.isDone(), q2.isDone())
	}
}

// TestScannerTiming verifies that ranges are scanned, regardless
// of how many, to match scanInterval.
func TestScannerTiming(t *testing.T) {
	defer leaktest.AfterTest(t)()
	const count = 3
	const runTime = 100 * time.Millisecond
	const maxError = 7500 * time.Microsecond
	durations := []time.Duration{
		15 * time.Millisecond,
		25 * time.Millisecond,
	}
	for i, duration := range durations {
		util.SucceedsSoon(t, func() error {
			ranges := newTestRangeSet(count, t)
			q := &testQueue{}
			s := newReplicaScanner(duration, 0, ranges)
			s.AddQueues(q)
			mc := hlc.NewManualClock(0)
			clock := hlc.NewClock(mc.UnixNano)
			stopper := stop.NewStopper()
			s.Start(clock, stopper)
			time.Sleep(runTime)
			stopper.Stop()

			avg := s.avgScan()
			log.Infof("%d: average scan: %s", i, avg)
			if avg.Nanoseconds()-duration.Nanoseconds() > maxError.Nanoseconds() ||
				duration.Nanoseconds()-avg.Nanoseconds() > maxError.Nanoseconds() {
				return util.Errorf("expected %s, got %s: exceeds max error of %s", duration, avg, maxError)
			}
			return nil
		})
	}
}

// TestScannerPaceInterval tests that paceInterval returns the correct interval.
func TestScannerPaceInterval(t *testing.T) {
	defer leaktest.AfterTest(t)()
	const count = 3
	durations := []time.Duration{
		30 * time.Millisecond,
		60 * time.Millisecond,
		500 * time.Millisecond,
	}
	// function logs an error when the actual value is not close
	// to the expected value
	logErrorWhenNotCloseTo := func(expected, actual time.Duration) {
		delta := 1 * time.Millisecond
		if actual < expected-delta || actual > expected+delta {
			t.Errorf("Expected duration %s, got %s", expected, actual)
		}
	}
	for _, duration := range durations {
		startTime := timeutil.Now()
		ranges := newTestRangeSet(count, t)
		s := newReplicaScanner(duration, 0, ranges)
		interval := s.paceInterval(startTime, startTime)
		logErrorWhenNotCloseTo(duration/count, interval)
		// The range set is empty
		ranges = newTestRangeSet(0, t)
		s = newReplicaScanner(duration, 0, ranges)
		interval = s.paceInterval(startTime, startTime)
		logErrorWhenNotCloseTo(duration, interval)
		ranges = newTestRangeSet(count, t)
		s = newReplicaScanner(duration, 0, ranges)
		// Move the present to duration time into the future
		interval = s.paceInterval(startTime, startTime.Add(duration))
		logErrorWhenNotCloseTo(0, interval)
	}
}

// TestScannerEmptyRangeSet verifies that an empty range set doesn't busy loop.
func TestScannerEmptyRangeSet(t *testing.T) {
	defer leaktest.AfterTest(t)()
	ranges := newTestRangeSet(0, t)
	q := &testQueue{}
	s := newReplicaScanner(time.Hour, 0, ranges)
	s.AddQueues(q)
	mc := hlc.NewManualClock(0)
	clock := hlc.NewClock(mc.UnixNano)
	stopper := stop.NewStopper()
	defer stopper.Stop()
	s.Start(clock, stopper)
	time.Sleep(time.Millisecond) // give it some time to (not) busy loop
	if count := s.Count(); count > 1 {
		t.Errorf("expected at most one loop, but got %d", count)
	}
}
<|endoftext|>