Switch to use upstream Prometheus interface

pkg/blockgen/blockgen.go

@@ -14,6 +14,7 @@ import (
 	"github.com/oklog/ulid"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/storage"
+	"github.com/prometheus/prometheus/tsdb/chunkenc"
 	"github.com/thanos-io/thanos/pkg/block/metadata"
 	"github.com/thanos-io/thanosbench/pkg/seriesgen"
 )
@@ -41,7 +42,7 @@ const (
 	Gauge   GenType = "GAUGE"
 )
 
-func (g GenType) Create(random *rand.Rand, mint, maxt int64, opts seriesgen.Characteristics) (seriesgen.SeriesIterator, error) {
+func (g GenType) Create(random *rand.Rand, mint, maxt int64, opts seriesgen.Characteristics) (chunkenc.Iterator, error) {
 	switch g {
 	case Random:
 		return seriesgen.NewValGen(random, mint, maxt, opts), nil
@@ -110,7 +111,7 @@ type blockSeriesSet struct {
 	target  int
 	err     error
 
-	curr seriesgen.Series
+	curr storage.Series
 }
 
 func (s *blockSeriesSet) Next() bool {
@@ -158,6 +159,8 @@ func (s *blockSeriesSet) Next() bool {
 	return true
 }
 
-func (s *blockSeriesSet) At() seriesgen.Series { return s.curr }
+func (s *blockSeriesSet) At() storage.Series { return s.curr }
 
 func (s *blockSeriesSet) Err() error { return s.err }
+
+func (s *blockSeriesSet) Warnings() storage.Warnings { return storage.Warnings{} }

pkg/seriesgen/append.go

@@ -8,15 +8,15 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
-func Append(ctx context.Context, goroutines int, appendable storage.Appendable, series SeriesSet) error {
+func Append(ctx context.Context, goroutines int, appendable storage.Appendable, series storage.SeriesSet) error {
 	g, gctx := errgroup.WithContext(ctx)
 
-	workBuffer := make(chan Series)
+	workBuffer := make(chan storage.Series)
 	for i := 0; i < goroutines; i++ {
 		app := appendable.Appender(gctx)
 		g.Go(func() error {
 			var (
-				s   Series
+				s   storage.Series
 				err error
 				ok  bool
 			)

pkg/seriesgen/append_test.go

@@ -19,7 +19,7 @@ import (
 
 type testSet struct {
 	count int
-	curr  Series
+	curr  storage.Series
 }
 
 func (s *testSet) Next() bool {
@@ -43,10 +43,12 @@ func (s *testSet) Next() bool {
 	return true
 }
 
-func (s *testSet) At() Series { return s.curr }
+func (s *testSet) At() storage.Series { return s.curr }
 
 func (s *testSet) Err() error { return nil }
 
+func (s *testSet) Warnings() storage.Warnings { return storage.Warnings{} }
+
 type testAppendable struct {
 	mtx     sync.Mutex
 	samples map[uint64][]sample

pkg/seriesgen/seriesgen.go

@@ -5,60 +5,34 @@ import (
 	"time"
 
 	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/storage"
+	"github.com/prometheus/prometheus/tsdb/chunkenc"
 )
 
 type sample struct {
 	T int64
 	V float64
 }
 
-// SeriesSet contains a set of series.
-type SeriesSet interface {
-	Next() bool
-	At() Series
-	Err() error
-}
-
-// Series exposes a single time series.
-type Series interface {
-	// Labels returns the complete set of labels identifying the series.
-	Labels() labels.Labels
-
-	// Iterator returns a new iterator of the data of the series.
-	// TODO(bwplotka): Consider moving this to tsdb.SeriesIterator if required. This means adding `Seek` method.
-	Iterator() SeriesIterator
-}
-
-// SeriesIterator iterates over the data of a time series.
-// Simplified version of github.com/prometheus/prometheus/tsdb.SeriesIterator
-type SeriesIterator interface {
-	// At returns the current timestamp/value pair.
-	At() (t int64, v float64)
-	// Next advances the iterator by one.
-	Next() bool
-	// Err returns current error.
-	Err() error
-}
-
 type SeriesGen struct {
-	SeriesIterator
+	chunkenc.Iterator
 
 	lset labels.Labels
 }
 
-func NewSeriesGen(lset labels.Labels, si SeriesIterator) *SeriesGen {
-	return &SeriesGen{
-		SeriesIterator: si,
-		lset:           lset,
+func NewSeriesGen(lset labels.Labels, it chunkenc.Iterator) storage.Series {
+	return &storage.SeriesEntry{
+		Lset: lset,
+		SampleIteratorFn: func() chunkenc.Iterator {
+			return it
+		},
 	}
 }
 func (s *SeriesGen) Labels() labels.Labels { return s.lset }
 
-func (s *SeriesGen) Iterator() SeriesIterator { return s }
-
-var _ SeriesIterator = &GaugeGen{}
-var _ SeriesIterator = &CounterGen{}
-var _ SeriesIterator = &ValGen{}
+var _ chunkenc.Iterator = &GaugeGen{}
+var _ chunkenc.Iterator = &CounterGen{}
+var _ chunkenc.Iterator = &ValGen{}
 
 type Characteristics struct {
 	Jitter         float64       `yaml:"jitter"`
@@ -124,6 +98,8 @@ func (g *GaugeGen) At() (t int64, v float64) {
 
 func (g *GaugeGen) Err() error { return nil }
 
+func (g *GaugeGen) Seek(_ int64) bool { return true }
+
 // TODO(bwplotka): Improve. Does not work well (: Too naive.
 // Add resets etc.
 type CounterGen struct {
@@ -215,6 +191,8 @@ func (g *CounterGen) At() (int64, float64) { return g.buff[0].T, g.buff[0].V }
 
 func (g *CounterGen) Err() error { return nil }
 
+func (g *CounterGen) Seek(_ int64) bool { return true }
+
 type ValGen struct {
 	interval         time.Duration
 	maxTime, minTime int64
@@ -252,3 +230,5 @@ func (g *ValGen) At() (t int64, v float64) {
 }
 
 func (g *ValGen) Err() error { return nil }
+
+func (g *ValGen) Seek(_ int64) bool { return true }

pkg/walgen/walgen.go

@@ -15,6 +15,7 @@ import (
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/model/timestamp"
+	"github.com/prometheus/prometheus/storage"
 	"github.com/prometheus/prometheus/tsdb"
 	"github.com/thanos-io/thanosbench/pkg/seriesgen"
 )
@@ -107,7 +108,7 @@ func GenerateTSDBWAL(logger log.Logger, dir string, config Config) error {
 }
 
 type Set struct {
-	s    []seriesgen.Series
+	s    []storage.Series
 	curr int
 }
 
@@ -119,6 +120,8 @@ func (s *Set) Next() bool {
 	return true
 }
 
-func (s *Set) At() seriesgen.Series { return s.s[s.curr-1] }
+func (s *Set) At() storage.Series { return s.s[s.curr-1] }
 
 func (s *Set) Err() error { return nil }
+
+func (s *Set) Warnings() storage.Warnings { return storage.Warnings{} }