github.com/influxdata/influxdb/v2@v2.7.6/tsdb/index_test.go (about)

     1  package tsdb_test
     2  
     3  import (
     4  	"compress/gzip"
     5  	"fmt"
     6  	"io"
     7  	"os"
     8  	"path/filepath"
     9  	"reflect"
    10  	"sync"
    11  	"testing"
    12  
    13  	"github.com/influxdata/influxdb/v2/influxql/query"
    14  	"github.com/influxdata/influxdb/v2/internal"
    15  	"github.com/influxdata/influxdb/v2/models"
    16  	"github.com/influxdata/influxdb/v2/pkg/slices"
    17  	"github.com/influxdata/influxdb/v2/tsdb"
    18  	"github.com/influxdata/influxdb/v2/tsdb/index/tsi1"
    19  	"github.com/influxdata/influxql"
    20  	"go.uber.org/zap/zaptest"
    21  )
    22  
    23  // Ensure iterator can merge multiple iterators together.
    24  func TestMergeSeriesIDIterators(t *testing.T) {
    25  	itr := tsdb.MergeSeriesIDIterators(
    26  		tsdb.NewSeriesIDSliceIterator([]uint64{1, 2, 3}),
    27  		tsdb.NewSeriesIDSliceIterator(nil),
    28  		nil,
    29  		tsdb.NewSeriesIDSliceIterator([]uint64{1, 2, 3, 4}),
    30  	)
    31  
    32  	if e, err := itr.Next(); err != nil {
    33  		t.Fatal(err)
    34  	} else if !reflect.DeepEqual(e, tsdb.SeriesIDElem{SeriesID: 1}) {
    35  		t.Fatalf("unexpected elem(0): %#v", e)
    36  	}
    37  	if e, err := itr.Next(); err != nil {
    38  		t.Fatal(err)
    39  	} else if !reflect.DeepEqual(e, tsdb.SeriesIDElem{SeriesID: 2}) {
    40  		t.Fatalf("unexpected elem(1): %#v", e)
    41  	}
    42  	if e, err := itr.Next(); err != nil {
    43  		t.Fatal(err)
    44  	} else if !reflect.DeepEqual(e, tsdb.SeriesIDElem{SeriesID: 3}) {
    45  		t.Fatalf("unexpected elem(2): %#v", e)
    46  	}
    47  	if e, err := itr.Next(); err != nil {
    48  		t.Fatal(err)
    49  	} else if !reflect.DeepEqual(e, tsdb.SeriesIDElem{SeriesID: 4}) {
    50  		t.Fatalf("unexpected elem(3): %#v", e)
    51  	}
    52  	if e, err := itr.Next(); err != nil {
    53  		t.Fatal(err)
    54  	} else if e.SeriesID != 0 {
    55  		t.Fatalf("expected nil elem: %#v", e)
    56  	}
    57  }
    58  
    59  func TestIndexSet_MeasurementNamesByExpr(t *testing.T) {
    60  	// Setup indexes
    61  	indexes := map[string]*Index{}
    62  	for _, name := range tsdb.RegisteredIndexes() {
    63  		idx := MustOpenNewIndex(t, name)
    64  		idx.AddSeries("cpu", map[string]string{"region": "east"})
    65  		idx.AddSeries("cpu", map[string]string{"region": "west", "secret": "foo"})
    66  		idx.AddSeries("disk", map[string]string{"secret": "foo"})
    67  		idx.AddSeries("mem", map[string]string{"region": "west"})
    68  		idx.AddSeries("gpu", map[string]string{"region": "east"})
    69  		idx.AddSeries("pci", map[string]string{"region": "east", "secret": "foo"})
    70  		indexes[name] = idx
    71  		defer idx.Close()
    72  	}
    73  
    74  	authorizer := &internal.AuthorizerMock{
    75  		AuthorizeSeriesReadFn: func(database string, measurement []byte, tags models.Tags) bool {
    76  			if tags.GetString("secret") != "" {
    77  				t.Logf("Rejecting series db=%s, m=%s, tags=%v", database, measurement, tags)
    78  				return false
    79  			}
    80  			return true
    81  		},
    82  	}
    83  
    84  	type example struct {
    85  		name     string
    86  		expr     influxql.Expr
    87  		expected [][]byte
    88  	}
    89  
    90  	// These examples should be run without any auth.
    91  	examples := []example{
    92  		{name: "all", expected: slices.StringsToBytes("cpu", "disk", "gpu", "mem", "pci")},
    93  		{name: "EQ", expr: influxql.MustParseExpr(`region = 'west'`), expected: slices.StringsToBytes("cpu", "mem")},
    94  		{name: "NEQ", expr: influxql.MustParseExpr(`region != 'west'`), expected: slices.StringsToBytes("gpu", "pci")},
    95  		{name: "EQREGEX", expr: influxql.MustParseExpr(`region =~ /.*st/`), expected: slices.StringsToBytes("cpu", "gpu", "mem", "pci")},
    96  		{name: "NEQREGEX", expr: influxql.MustParseExpr(`region !~ /.*est/`), expected: slices.StringsToBytes("gpu", "pci")},
    97  	}
    98  
    99  	// These examples should be run with the authorizer.
   100  	authExamples := []example{
   101  		{name: "all", expected: slices.StringsToBytes("cpu", "gpu", "mem")},
   102  		{name: "EQ", expr: influxql.MustParseExpr(`region = 'west'`), expected: slices.StringsToBytes("mem")},
   103  		{name: "NEQ", expr: influxql.MustParseExpr(`region != 'west'`), expected: slices.StringsToBytes("gpu")},
   104  		{name: "EQREGEX", expr: influxql.MustParseExpr(`region =~ /.*st/`), expected: slices.StringsToBytes("cpu", "gpu", "mem")},
   105  		{name: "NEQREGEX", expr: influxql.MustParseExpr(`region !~ /.*est/`), expected: slices.StringsToBytes("gpu")},
   106  	}
   107  
   108  	for _, idx := range tsdb.RegisteredIndexes() {
   109  		t.Run(idx, func(t *testing.T) {
   110  			t.Run("no authorization", func(t *testing.T) {
   111  				for _, example := range examples {
   112  					t.Run(example.name, func(t *testing.T) {
   113  						names, err := indexes[idx].IndexSet().MeasurementNamesByExpr(nil, example.expr)
   114  						if err != nil {
   115  							t.Fatal(err)
   116  						} else if !reflect.DeepEqual(names, example.expected) {
   117  							t.Fatalf("got names: %v, expected %v", slices.BytesToStrings(names), slices.BytesToStrings(example.expected))
   118  						}
   119  					})
   120  				}
   121  			})
   122  
   123  			t.Run("with authorization", func(t *testing.T) {
   124  				for _, example := range authExamples {
   125  					t.Run(example.name, func(t *testing.T) {
   126  						names, err := indexes[idx].IndexSet().MeasurementNamesByExpr(authorizer, example.expr)
   127  						if err != nil {
   128  							t.Fatal(err)
   129  						} else if !reflect.DeepEqual(names, example.expected) {
   130  							t.Fatalf("got names: %v, expected %v", slices.BytesToStrings(names), slices.BytesToStrings(example.expected))
   131  						}
   132  					})
   133  				}
   134  			})
   135  		})
   136  	}
   137  }
   138  
   139  func TestIndexSet_MeasurementNamesByPredicate(t *testing.T) {
   140  	// Setup indexes
   141  	indexes := map[string]*Index{}
   142  	for _, name := range tsdb.RegisteredIndexes() {
   143  		idx := MustOpenNewIndex(t, name)
   144  		idx.AddSeries("cpu", map[string]string{"region": "east"})
   145  		idx.AddSeries("cpu", map[string]string{"region": "west", "secret": "foo"})
   146  		idx.AddSeries("disk", map[string]string{"secret": "foo"})
   147  		idx.AddSeries("mem", map[string]string{"region": "west"})
   148  		idx.AddSeries("gpu", map[string]string{"region": "east"})
   149  		idx.AddSeries("pci", map[string]string{"region": "east", "secret": "foo"})
   150  		indexes[name] = idx
   151  		defer idx.Close()
   152  	}
   153  
   154  	authorizer := &internal.AuthorizerMock{
   155  		AuthorizeSeriesReadFn: func(database string, measurement []byte, tags models.Tags) bool {
   156  			if tags.GetString("secret") != "" {
   157  				t.Logf("Rejecting series db=%s, m=%s, tags=%v", database, measurement, tags)
   158  				return false
   159  			}
   160  			return true
   161  		},
   162  	}
   163  
   164  	type example struct {
   165  		name     string
   166  		expr     influxql.Expr
   167  		expected [][]byte
   168  	}
   169  
   170  	// These examples should be run without any auth.
   171  	examples := []example{
   172  		{name: "all", expected: slices.StringsToBytes("cpu", "disk", "gpu", "mem", "pci")},
   173  		{name: "EQ", expr: influxql.MustParseExpr(`region = 'west'`), expected: slices.StringsToBytes("cpu", "mem")},
   174  		{name: "NEQ", expr: influxql.MustParseExpr(`region != 'west'`), expected: slices.StringsToBytes("cpu", "disk", "gpu", "pci")},
   175  		{name: "EQREGEX", expr: influxql.MustParseExpr(`region =~ /.*st/`), expected: slices.StringsToBytes("cpu", "gpu", "mem", "pci")},
   176  		{name: "NEQREGEX", expr: influxql.MustParseExpr(`region !~ /.*est/`), expected: slices.StringsToBytes("cpu", "disk", "gpu", "pci")},
   177  		// None of the series have this tag so all should be selected.
   178  		{name: "EQ empty", expr: influxql.MustParseExpr(`host = ''`), expected: slices.StringsToBytes("cpu", "disk", "gpu", "mem", "pci")},
   179  		// Measurements that have this tag at all should be returned.
   180  		{name: "NEQ empty", expr: influxql.MustParseExpr(`region != ''`), expected: slices.StringsToBytes("cpu", "gpu", "mem", "pci")},
   181  		{name: "EQREGEX empty", expr: influxql.MustParseExpr(`host =~ /.*/`), expected: slices.StringsToBytes("cpu", "disk", "gpu", "mem", "pci")},
   182  		{name: "NEQ empty", expr: influxql.MustParseExpr(`region !~ /.*/`), expected: slices.StringsToBytes()},
   183  	}
   184  
   185  	// These examples should be run with the authorizer.
   186  	authExamples := []example{
   187  		{name: "all", expected: slices.StringsToBytes("cpu", "gpu", "mem")},
   188  		{name: "EQ", expr: influxql.MustParseExpr(`region = 'west'`), expected: slices.StringsToBytes("mem")},
   189  		{name: "NEQ", expr: influxql.MustParseExpr(`region != 'west'`), expected: slices.StringsToBytes("cpu", "gpu")},
   190  		{name: "EQREGEX", expr: influxql.MustParseExpr(`region =~ /.*st/`), expected: slices.StringsToBytes("cpu", "gpu", "mem")},
   191  		{name: "NEQREGEX", expr: influxql.MustParseExpr(`region !~ /.*est/`), expected: slices.StringsToBytes("cpu", "gpu")},
   192  		{name: "EQ empty", expr: influxql.MustParseExpr(`host = ''`), expected: slices.StringsToBytes("cpu", "gpu", "mem")},
   193  		{name: "NEQ empty", expr: influxql.MustParseExpr(`region != ''`), expected: slices.StringsToBytes("cpu", "gpu", "mem")},
   194  		{name: "EQREGEX empty", expr: influxql.MustParseExpr(`host =~ /.*/`), expected: slices.StringsToBytes("cpu", "gpu", "mem")},
   195  		{name: "NEQ empty", expr: influxql.MustParseExpr(`region !~ /.*/`), expected: slices.StringsToBytes()},
   196  	}
   197  
   198  	for _, idx := range tsdb.RegisteredIndexes() {
   199  		t.Run(idx, func(t *testing.T) {
   200  			t.Run("no authorization", func(t *testing.T) {
   201  				for _, example := range examples {
   202  					t.Run(example.name, func(t *testing.T) {
   203  						names, err := indexes[idx].IndexSet().MeasurementNamesByPredicate(nil, example.expr)
   204  						if err != nil {
   205  							t.Fatal(err)
   206  						} else if !reflect.DeepEqual(names, example.expected) {
   207  							t.Fatalf("got names: %v, expected %v", slices.BytesToStrings(names), slices.BytesToStrings(example.expected))
   208  						}
   209  					})
   210  				}
   211  			})
   212  
   213  			t.Run("with authorization", func(t *testing.T) {
   214  				for _, example := range authExamples {
   215  					t.Run(example.name, func(t *testing.T) {
   216  						names, err := indexes[idx].IndexSet().MeasurementNamesByPredicate(authorizer, example.expr)
   217  						if err != nil {
   218  							t.Fatal(err)
   219  						} else if !reflect.DeepEqual(names, example.expected) {
   220  							t.Fatalf("got names: %v, expected %v", slices.BytesToStrings(names), slices.BytesToStrings(example.expected))
   221  						}
   222  					})
   223  				}
   224  			})
   225  		})
   226  	}
   227  }
   228  
   229  func TestIndex_Sketches(t *testing.T) {
   230  	checkCardinalities := func(t *testing.T, index *Index, state string, series, tseries, measurements, tmeasurements int) {
   231  		t.Helper()
   232  
   233  		// Get sketches and check cardinality...
   234  		sketch, tsketch, err := index.SeriesSketches()
   235  		if err != nil {
   236  			t.Fatal(err)
   237  		}
   238  
   239  		// delta calculates a rough 10% delta. If i is small then a minimum value
   240  		// of 2 is used.
   241  		delta := func(i int) int {
   242  			v := i / 10
   243  			if v == 0 {
   244  				v = 2
   245  			}
   246  			return v
   247  		}
   248  
   249  		// series cardinality should be well within 10%.
   250  		if got, exp := int(sketch.Count()), series; got-exp < -delta(series) || got-exp > delta(series) {
   251  			t.Errorf("[%s] got series cardinality %d, expected ~%d", state, got, exp)
   252  		}
   253  
   254  		// check series tombstones
   255  		if got, exp := int(tsketch.Count()), tseries; got-exp < -delta(tseries) || got-exp > delta(tseries) {
   256  			t.Errorf("[%s] got series tombstone cardinality %d, expected ~%d", state, got, exp)
   257  		}
   258  
   259  		// Check measurement cardinality.
   260  		if sketch, tsketch, err = index.MeasurementsSketches(); err != nil {
   261  			t.Fatal(err)
   262  		}
   263  
   264  		if got, exp := int(sketch.Count()), measurements; got != exp { //got-exp < -delta(measurements) || got-exp > delta(measurements) {
   265  			t.Errorf("[%s] got measurement cardinality %d, expected ~%d", state, got, exp)
   266  		}
   267  
   268  		if got, exp := int(tsketch.Count()), tmeasurements; got != exp { //got-exp < -delta(tmeasurements) || got-exp > delta(tmeasurements) {
   269  			t.Errorf("[%s] got measurement tombstone cardinality %d, expected ~%d", state, got, exp)
   270  		}
   271  	}
   272  
   273  	test := func(t *testing.T, index string) error {
   274  		idx := MustNewIndex(t, index)
   275  		if index, ok := idx.Index.(*tsi1.Index); ok {
   276  			// Override the log file max size to force a log file compaction sooner.
   277  			// This way, we will test the sketches are correct when they have been
   278  			// compacted into IndexFiles, and also when they're loaded from
   279  			// IndexFiles after a re-open.
   280  			tsi1.WithMaximumLogFileSize(1 << 10)(index)
   281  		}
   282  
   283  		// Open the index
   284  		idx.MustOpen()
   285  		defer idx.Close()
   286  
   287  		series := genTestSeries(10, 5, 3)
   288  		// Add series to index.
   289  		for _, serie := range series {
   290  			if err := idx.AddSeries(serie.Measurement, serie.Tags.Map()); err != nil {
   291  				t.Fatal(err)
   292  			}
   293  		}
   294  
   295  		// Check cardinalities after adding series.
   296  		checkCardinalities(t, idx, "initial", 2430, 0, 10, 0)
   297  
   298  		// Re-open step only applies to the TSI index.
   299  		if _, ok := idx.Index.(*tsi1.Index); ok {
   300  			// Re-open the index.
   301  			if err := idx.Reopen(); err != nil {
   302  				panic(err)
   303  			}
   304  
   305  			// Check cardinalities after the reopen
   306  			checkCardinalities(t, idx, "initial|reopen", 2430, 0, 10, 0)
   307  		}
   308  
   309  		// Drop some series
   310  		if err := idx.DropMeasurement([]byte("measurement2")); err != nil {
   311  			return err
   312  		} else if err := idx.DropMeasurement([]byte("measurement5")); err != nil {
   313  			return err
   314  		}
   315  
   316  		// Check cardinalities after the delete
   317  		checkCardinalities(t, idx, "initial|reopen|delete", 2430, 486, 10, 2)
   318  
   319  		// Re-open step only applies to the TSI index.
   320  		if _, ok := idx.Index.(*tsi1.Index); ok {
   321  			// Re-open the index.
   322  			if err := idx.Reopen(); err != nil {
   323  				panic(err)
   324  			}
   325  
   326  			// Check cardinalities after the reopen
   327  			checkCardinalities(t, idx, "initial|reopen|delete|reopen", 2430, 486, 10, 2)
   328  		}
   329  		return nil
   330  	}
   331  
   332  	for _, index := range tsdb.RegisteredIndexes() {
   333  		t.Run(index, func(t *testing.T) {
   334  			if err := test(t, index); err != nil {
   335  				t.Fatal(err)
   336  			}
   337  		})
   338  	}
   339  }
   340  
   341  // Index wraps a series file and index.
   342  type Index struct {
   343  	tsdb.Index
   344  	rootPath  string
   345  	indexType string
   346  	sfile     *tsdb.SeriesFile
   347  }
   348  
   349  type EngineOption func(opts *tsdb.EngineOptions)
   350  
   351  // DisableTSICache allows the caller to disable the TSI bitset cache during a test.
   352  var DisableTSICache = func() EngineOption {
   353  	return func(opts *tsdb.EngineOptions) {
   354  		opts.Config.SeriesIDSetCacheSize = 0
   355  	}
   356  }
   357  
   358  // MustNewIndex will initialize a new index using the provide type. It creates
   359  // everything under the same root directory so it can be cleanly removed on Close.
   360  //
   361  // The index will not be opened.
   362  func MustNewIndex(tb testing.TB, index string, eopts ...EngineOption) *Index {
   363  	tb.Helper()
   364  
   365  	opts := tsdb.NewEngineOptions()
   366  	opts.IndexVersion = index
   367  
   368  	for _, opt := range eopts {
   369  		opt(&opts)
   370  	}
   371  
   372  	rootPath := tb.TempDir()
   373  
   374  	seriesPath, err := os.MkdirTemp(rootPath, tsdb.SeriesFileDirectory)
   375  	if err != nil {
   376  		panic(err)
   377  	}
   378  
   379  	sfile := tsdb.NewSeriesFile(seriesPath)
   380  	if err := sfile.Open(); err != nil {
   381  		panic(err)
   382  	}
   383  
   384  	i, err := tsdb.NewIndex(0, "db0", filepath.Join(rootPath, "index"), tsdb.NewSeriesIDSet(), sfile, opts)
   385  	if err != nil {
   386  		panic(err)
   387  	}
   388  	i.WithLogger(zaptest.NewLogger(tb))
   389  
   390  	idx := &Index{
   391  		Index:     i,
   392  		indexType: index,
   393  		rootPath:  rootPath,
   394  		sfile:     sfile,
   395  	}
   396  	return idx
   397  }
   398  
   399  // MustOpenNewIndex will initialize a new index using the provide type and opens
   400  // it.
   401  func MustOpenNewIndex(tb testing.TB, index string, opts ...EngineOption) *Index {
   402  	tb.Helper()
   403  
   404  	idx := MustNewIndex(tb, index, opts...)
   405  	idx.MustOpen()
   406  	return idx
   407  }
   408  
   409  // MustOpen opens the underlying index or panics.
   410  func (i *Index) MustOpen() {
   411  	if err := i.Index.Open(); err != nil {
   412  		panic(err)
   413  	}
   414  }
   415  
   416  func (idx *Index) IndexSet() *tsdb.IndexSet {
   417  	return &tsdb.IndexSet{Indexes: []tsdb.Index{idx.Index}, SeriesFile: idx.sfile}
   418  }
   419  
   420  func (idx *Index) AddSeries(name string, tags map[string]string) error {
   421  	t := models.NewTags(tags)
   422  	key := fmt.Sprintf("%s,%s", name, t.HashKey())
   423  	return idx.CreateSeriesIfNotExists([]byte(key), []byte(name), t)
   424  }
   425  
   426  // Reopen closes and re-opens the underlying index, without removing any data.
   427  func (i *Index) Reopen() error {
   428  	if err := i.Index.Close(); err != nil {
   429  		return err
   430  	}
   431  
   432  	if err := i.sfile.Close(); err != nil {
   433  		return err
   434  	}
   435  
   436  	i.sfile = tsdb.NewSeriesFile(i.sfile.Path())
   437  	if err := i.sfile.Open(); err != nil {
   438  		return err
   439  	}
   440  
   441  	opts := tsdb.NewEngineOptions()
   442  	opts.IndexVersion = i.indexType
   443  
   444  	idx, err := tsdb.NewIndex(0, "db0", filepath.Join(i.rootPath, "index"), tsdb.NewSeriesIDSet(), i.sfile, opts)
   445  	if err != nil {
   446  		return err
   447  	}
   448  	i.Index = idx
   449  	return i.Index.Open()
   450  }
   451  
   452  // Close closes the index cleanly and removes all on-disk data.
   453  func (i *Index) Close() error {
   454  	if err := i.Index.Close(); err != nil {
   455  		return err
   456  	}
   457  
   458  	if err := i.sfile.Close(); err != nil {
   459  		return err
   460  	}
   461  	//return os.RemoveAll(i.rootPath)
   462  	return nil
   463  }
   464  
   465  // This benchmark compares the TagSets implementation across index types.
   466  //
   467  // In the case of the TSI index, TagSets has to merge results across all several
   468  // index partitions.
   469  //
   470  // Typical results on an i7 laptop.
   471  //
   472  // BenchmarkIndexSet_TagSets/1M_series/tsi1-8    	     100	  18995530 ns/op	 5221180 B/op	   20379 allocs/op
   473  func BenchmarkIndexSet_TagSets(b *testing.B) {
   474  	// Read line-protocol and coerce into tsdb format.
   475  	keys := make([][]byte, 0, 1e6)
   476  	names := make([][]byte, 0, 1e6)
   477  	tags := make([]models.Tags, 0, 1e6)
   478  
   479  	// 1M series generated with:
   480  	// $inch -b 10000 -c 1 -t 10,10,10,10,10,10 -f 1 -m 5 -p 1
   481  	fd, err := os.Open("testdata/line-protocol-1M.txt.gz")
   482  	if err != nil {
   483  		b.Fatal(err)
   484  	}
   485  
   486  	gzr, err := gzip.NewReader(fd)
   487  	if err != nil {
   488  		fd.Close()
   489  		b.Fatal(err)
   490  	}
   491  
   492  	data, err := io.ReadAll(gzr)
   493  	if err != nil {
   494  		b.Fatal(err)
   495  	}
   496  
   497  	if err := fd.Close(); err != nil {
   498  		b.Fatal(err)
   499  	}
   500  
   501  	points, err := models.ParsePoints(data)
   502  	if err != nil {
   503  		b.Fatal(err)
   504  	}
   505  
   506  	for _, pt := range points {
   507  		keys = append(keys, pt.Key())
   508  		names = append(names, pt.Name())
   509  		tags = append(tags, pt.Tags())
   510  	}
   511  
   512  	// setup writes all of the above points to the index.
   513  	setup := func(idx *Index) {
   514  		batchSize := 10000
   515  		for j := 0; j < 1; j++ {
   516  			for i := 0; i < len(keys); i += batchSize {
   517  				k := keys[i : i+batchSize]
   518  				n := names[i : i+batchSize]
   519  				t := tags[i : i+batchSize]
   520  				if err := idx.CreateSeriesListIfNotExists(k, n, t); err != nil {
   521  					b.Fatal(err)
   522  				}
   523  			}
   524  		}
   525  	}
   526  
   527  	var errResult error
   528  
   529  	// This benchmark will merge eight bitsets each containing ~10,000 series IDs.
   530  	b.Run("1M series", func(b *testing.B) {
   531  		b.ReportAllocs()
   532  		for _, indexType := range tsdb.RegisteredIndexes() {
   533  			idx := MustOpenNewIndex(b, indexType)
   534  			setup(idx)
   535  
   536  			name := []byte("m4")
   537  			opt := query.IteratorOptions{Condition: influxql.MustParseExpr(`"tag5"::tag = 'value0'`)}
   538  			indexSet := tsdb.IndexSet{
   539  				SeriesFile: idx.sfile,
   540  				Indexes:    []tsdb.Index{idx.Index},
   541  			} // For TSI implementation
   542  
   543  			ts := func() ([]*query.TagSet, error) {
   544  				return indexSet.TagSets(idx.sfile, name, opt)
   545  			}
   546  
   547  			b.Run(indexType, func(b *testing.B) {
   548  				for i := 0; i < b.N; i++ {
   549  					// Will call TagSets on the appropriate implementation.
   550  					_, errResult = ts()
   551  					if errResult != nil {
   552  						b.Fatal(err)
   553  					}
   554  				}
   555  			})
   556  
   557  			if err := idx.Close(); err != nil {
   558  				b.Fatal(err)
   559  			}
   560  		}
   561  	})
   562  }
   563  
   564  // This benchmark concurrently writes series to the index and fetches cached bitsets.
   565  // The idea is to emphasize the performance difference when bitset caching is on and off.
   566  //
   567  // # Typical results for an i7 laptop
   568  //
   569  // BenchmarkIndex_ConcurrentWriteQuery/tsi1/queries_100000/cache-8        1	1645048376 ns/op	2215402840 B/op	 23048978 allocs/op
   570  // BenchmarkIndex_ConcurrentWriteQuery/tsi1/queries_100000/no_cache-8     1	22242155616 ns/op	28277544136 B/op 79620463 allocs/op
   571  func BenchmarkIndex_ConcurrentWriteQuery(b *testing.B) {
   572  	// Read line-protocol and coerce into tsdb format.
   573  	keys := make([][]byte, 0, 1e6)
   574  	names := make([][]byte, 0, 1e6)
   575  	tags := make([]models.Tags, 0, 1e6)
   576  
   577  	// 1M series generated with:
   578  	// $inch -b 10000 -c 1 -t 10,10,10,10,10,10 -f 1 -m 5 -p 1
   579  	fd, err := os.Open("testdata/line-protocol-1M.txt.gz")
   580  	if err != nil {
   581  		b.Fatal(err)
   582  	}
   583  
   584  	gzr, err := gzip.NewReader(fd)
   585  	if err != nil {
   586  		fd.Close()
   587  		b.Fatal(err)
   588  	}
   589  
   590  	data, err := io.ReadAll(gzr)
   591  	if err != nil {
   592  		b.Fatal(err)
   593  	}
   594  
   595  	if err := fd.Close(); err != nil {
   596  		b.Fatal(err)
   597  	}
   598  
   599  	points, err := models.ParsePoints(data)
   600  	if err != nil {
   601  		b.Fatal(err)
   602  	}
   603  
   604  	for _, pt := range points {
   605  		keys = append(keys, pt.Key())
   606  		names = append(names, pt.Name())
   607  		tags = append(tags, pt.Tags())
   608  	}
   609  
   610  	runBenchmark := func(b *testing.B, index string, queryN int, useTSICache bool) {
   611  		var idx *Index
   612  		if !useTSICache {
   613  			idx = MustOpenNewIndex(b, index, DisableTSICache())
   614  		} else {
   615  			idx = MustOpenNewIndex(b, index)
   616  		}
   617  
   618  		var wg sync.WaitGroup
   619  		begin := make(chan struct{})
   620  
   621  		// Run concurrent iterator...
   622  		runIter := func() {
   623  			keys := [][]string{
   624  				{"m0", "tag2", "value4"},
   625  				{"m1", "tag3", "value5"},
   626  				{"m2", "tag4", "value6"},
   627  				{"m3", "tag0", "value8"},
   628  				{"m4", "tag5", "value0"},
   629  			}
   630  
   631  			<-begin // Wait for writes to land
   632  			for i := 0; i < queryN/5; i++ {
   633  				for _, key := range keys {
   634  					itr, err := idx.TagValueSeriesIDIterator([]byte(key[0]), []byte(key[1]), []byte(key[2]))
   635  					if err != nil {
   636  						b.Fatal(err)
   637  					}
   638  
   639  					if itr == nil {
   640  						panic("should not happen")
   641  					}
   642  
   643  					if err := itr.Close(); err != nil {
   644  						b.Fatal(err)
   645  					}
   646  				}
   647  			}
   648  		}
   649  
   650  		batchSize := 10000
   651  		wg.Add(1)
   652  		go func() { defer wg.Done(); runIter() }()
   653  		var once sync.Once
   654  		for j := 0; j < b.N; j++ {
   655  			for i := 0; i < len(keys); i += batchSize {
   656  				k := keys[i : i+batchSize]
   657  				n := names[i : i+batchSize]
   658  				t := tags[i : i+batchSize]
   659  				if err := idx.CreateSeriesListIfNotExists(k, n, t); err != nil {
   660  					b.Fatal(err)
   661  				}
   662  				once.Do(func() { close(begin) })
   663  			}
   664  
   665  			// Wait for queries to finish
   666  			wg.Wait()
   667  
   668  			// Reset the index...
   669  			b.StopTimer()
   670  			if err := idx.Close(); err != nil {
   671  				b.Fatal(err)
   672  			}
   673  
   674  			// Re-open everything
   675  			idx = MustOpenNewIndex(b, index)
   676  			wg.Add(1)
   677  			begin = make(chan struct{})
   678  			once = sync.Once{}
   679  			go func() { defer wg.Done(); runIter() }()
   680  			b.StartTimer()
   681  		}
   682  	}
   683  
   684  	queries := []int{1e5}
   685  	for _, indexType := range tsdb.RegisteredIndexes() {
   686  		b.Run(indexType, func(b *testing.B) {
   687  			for _, queryN := range queries {
   688  				b.Run(fmt.Sprintf("queries %d", queryN), func(b *testing.B) {
   689  					b.Run("cache", func(b *testing.B) {
   690  						runBenchmark(b, indexType, queryN, true)
   691  					})
   692  
   693  					b.Run("no cache", func(b *testing.B) {
   694  						runBenchmark(b, indexType, queryN, false)
   695  					})
   696  				})
   697  			}
   698  		})
   699  	}
   700  }