package tsdb import ( "math/rand" "sort" "testing" "github.com/fabxc/tsdb/labels" "github.com/stretchr/testify/require" ) type mockSeriesIterator struct { seek func(int64) bool at func() (int64, float64) next func() bool err func() error } func (m *mockSeriesIterator) Seek(t int64) bool { return m.seek(t) } func (m *mockSeriesIterator) At() (int64, float64) { return m.at() } func (m *mockSeriesIterator) Next() bool { return m.next() } func (m *mockSeriesIterator) Err() error { return m.err() } type mockSeries struct { labels func() labels.Labels iterator func() SeriesIterator } func (m *mockSeries) Labels() labels.Labels { return m.labels() } func (m *mockSeries) Iterator() SeriesIterator { return m.iterator() } type listSeriesIterator struct { list []sample idx int } func newListSeriesIterator(list []sample) *listSeriesIterator { return &listSeriesIterator{list: list, idx: -1} } func (it *listSeriesIterator) At() (int64, float64) { s := it.list[it.idx] return s.t, s.v } func (it *listSeriesIterator) Next() bool { it.idx++ return it.idx < len(it.list) } func (it *listSeriesIterator) Seek(t int64) bool { if it.idx == -1 { it.idx = 0 } // Do binary search between current position and end. it.idx = sort.Search(len(it.list)-it.idx, func(i int) bool { s := it.list[i+it.idx] return s.t >= t }) return it.idx < len(it.list) } func (it *listSeriesIterator) Err() error { return nil } type mockSeriesSet struct { next func() bool series func() Series err func() error } func (m *mockSeriesSet) Next() bool { return m.next() } func (m *mockSeriesSet) At() Series { return m.series() } func (m *mockSeriesSet) Err() error { return m.err() } func newListSeriesSet(list []Series) *mockSeriesSet { i := -1 return &mockSeriesSet{ next: func() bool { i++ return i < len(list) }, series: func() Series { return list[i] }, err: func() error { return nil }, } } func TestShardSeriesSet(t *testing.T) { newSeries := func(l map[string]string, s []sample) Series { return &mockSeries{ labels: func() labels.Labels { return labels.FromMap(l) }, iterator: func() SeriesIterator { return newListSeriesIterator(s) }, } } cases := []struct { // The input sets in order (samples in series in b are strictly // after those in a). a, b SeriesSet // The composition of a and b in the shard series set must yield // results equivalent to the result series set. exp SeriesSet }{ { a: newListSeriesSet([]Series{ newSeries(map[string]string{ "a": "a", }, []sample{ {t: 1, v: 1}, }), }), b: newListSeriesSet([]Series{ newSeries(map[string]string{ "a": "a", }, []sample{ {t: 2, v: 2}, }), newSeries(map[string]string{ "b": "b", }, []sample{ {t: 1, v: 1}, }), }), exp: newListSeriesSet([]Series{ newSeries(map[string]string{ "a": "a", }, []sample{ {t: 1, v: 1}, {t: 2, v: 2}, }), newSeries(map[string]string{ "b": "b", }, []sample{ {t: 1, v: 1}, }), }), }, } Outer: for _, c := range cases { res := newShardSeriesSet(c.a, c.b) for { eok, rok := c.exp.Next(), res.Next() require.Equal(t, eok, rok, "next") if !eok { continue Outer } sexp := c.exp.At() sres := res.At() require.Equal(t, sexp.Labels(), sres.Labels(), "labels") smplExp, errExp := expandSeriesIterator(sexp.Iterator()) smplRes, errRes := expandSeriesIterator(sres.Iterator()) require.Equal(t, errExp, errRes, "samples error") require.Equal(t, smplExp, smplRes, "samples") } } } func expandSeriesIterator(it SeriesIterator) (r []sample, err error) { for it.Next() { t, v := it.At() r = append(r, sample{t: t, v: v}) } return r, it.Err() } func TestSampleRing(t *testing.T) { cases := []struct { input []int64 delta int64 size int }{ { input: []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, delta: 2, size: 1, }, { input: []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, delta: 2, size: 2, }, { input: []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, delta: 7, size: 3, }, { input: []int64{1, 2, 3, 4, 5, 16, 17, 18, 19, 20}, delta: 7, size: 1, }, } for _, c := range cases { r := newSampleRing(c.delta, c.size) input := []sample{} for _, t := range c.input { input = append(input, sample{ t: t, v: float64(rand.Intn(100)), }) } for i, s := range input { r.add(s.t, s.v) buffered := r.samples() for _, sold := range input[:i] { found := false for _, bs := range buffered { if bs.t == sold.t && bs.v == sold.v { found = true break } } if sold.t >= s.t-c.delta && !found { t.Fatalf("%d: expected sample %d to be in buffer but was not; buffer %v", i, sold.t, buffered) } if sold.t < s.t-c.delta && found { t.Fatalf("%d: unexpected sample %d in buffer; buffer %v", i, sold.t, buffered) } } } } } func TestBufferedSeriesIterator(t *testing.T) { var it *BufferedSeriesIterator bufferEq := func(exp []sample) { var b []sample bit := it.Buffer() for bit.Next() { t, v := bit.At() b = append(b, sample{t: t, v: v}) } require.Equal(t, exp, b, "buffer mismatch") } sampleEq := func(ets int64, ev float64) { ts, v := it.At() require.Equal(t, ets, ts, "timestamp mismatch") require.Equal(t, ev, v, "value mismatch") } it = NewBuffer(newListSeriesIterator([]sample{ {t: 1, v: 2}, {t: 2, v: 3}, {t: 3, v: 4}, {t: 4, v: 5}, {t: 5, v: 6}, {t: 99, v: 8}, {t: 100, v: 9}, {t: 101, v: 10}, }), 2) require.True(t, it.Seek(-123), "seek failed") sampleEq(1, 2) bufferEq(nil) require.True(t, it.Next(), "next failed") sampleEq(2, 3) bufferEq([]sample{{t: 1, v: 2}}) require.True(t, it.Next(), "next failed") require.True(t, it.Next(), "next failed") require.True(t, it.Next(), "next failed") sampleEq(5, 6) bufferEq([]sample{{t: 2, v: 3}, {t: 3, v: 4}, {t: 4, v: 5}}) require.True(t, it.Seek(5), "seek failed") sampleEq(5, 6) bufferEq([]sample{{t: 2, v: 3}, {t: 3, v: 4}, {t: 4, v: 5}}) require.True(t, it.Seek(101), "seek failed") sampleEq(101, 10) bufferEq([]sample{{t: 99, v: 8}, {t: 100, v: 9}}) require.False(t, it.Next(), "next succeeded unexpectedly") }