github.com/grafana/pyroscope@v1.18.0/pkg/metastore/index/index_bench_test.go (about)

     1  package index
     2  
     3  import (
     4  	"fmt"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/stretchr/testify/require"
     9  	"go.etcd.io/bbolt"
    10  
    11  	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
    12  	"github.com/grafana/pyroscope/pkg/test"
    13  	"github.com/grafana/pyroscope/pkg/util"
    14  )
    15  
    16  func BenchmarkIndex_GetTenantStats(b *testing.B) {
    17  	const (
    18  		partitionDuration = 6 * time.Hour
    19  		numPartitions     = 6 * 4 * 30
    20  		numTenants        = 100
    21  		numShards         = 1
    22  	)
    23  
    24  	db := test.BoltDB(b)
    25  	defer func() {
    26  		require.NoError(b, db.Close())
    27  	}()
    28  
    29  	config := DefaultConfig
    30  	config.partitionDuration = partitionDuration
    31  	config.ShardCacheSize = 1000
    32  	config.BlockReadCacheSize = 1000
    33  	config.BlockWriteCacheSize = 1000
    34  
    35  	idx := NewIndex(util.Logger, NewStore(), config)
    36  	require.NoError(b, db.Update(idx.Init))
    37  
    38  	startTime := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
    39  	var blocks int
    40  
    41  	for i := 0; i < numPartitions; i++ {
    42  		start := startTime.Add(time.Duration(i) * partitionDuration)
    43  		for j := 0; j < numTenants; j++ {
    44  			tenant := fmt.Sprintf("tenant-%03d", j)
    45  			for shard := 0; shard < numShards; shard++ {
    46  				ts := start.Add(time.Duration(shard) * time.Hour)
    47  				md := &metastorev1.BlockMeta{
    48  					FormatVersion: 1,
    49  					Id:            test.ULID(ts.Format(time.RFC3339)),
    50  					Tenant:        1,
    51  					Shard:         uint32(shard),
    52  					MinTime:       ts.UnixMilli(),
    53  					MaxTime:       ts.Add(partitionDuration / 2).UnixMilli(),
    54  					StringTable:   []string{"", tenant},
    55  				}
    56  				err := db.Update(func(tx *bbolt.Tx) error {
    57  					return idx.InsertBlock(tx, md)
    58  				})
    59  				require.NoError(b, err)
    60  				blocks++
    61  			}
    62  		}
    63  
    64  		if (i+1)%100 == 0 {
    65  			b.Logf("Created %d/%d partitions (%d blocks so far)", i+1, numPartitions, blocks)
    66  		}
    67  	}
    68  
    69  	for _, tc := range []struct {
    70  		name   string
    71  		tenant string
    72  		desc   string
    73  	}{
    74  		{
    75  			name:   "ExistingTenant",
    76  			tenant: "tenant-000",
    77  			desc:   "Tenant with data in all partitions",
    78  		},
    79  		{
    80  			name:   "MidTenant",
    81  			tenant: "tenant-050",
    82  			desc:   "Tenant in the middle range",
    83  		},
    84  		{
    85  			name:   "LastTenant",
    86  			tenant: "tenant-099",
    87  			desc:   "Last tenant in the range",
    88  		},
    89  		{
    90  			name:   "NonExistentTenant",
    91  			tenant: "tenant-999",
    92  			desc:   "Tenant that doesn't exist",
    93  		},
    94  	} {
    95  		b.Run(tc.name, func(b *testing.B) {
    96  			b.ResetTimer()
    97  
    98  			for i := 0; i < b.N; i++ {
    99  				var stats *metastorev1.TenantStats
   100  				err := db.View(func(tx *bbolt.Tx) error {
   101  					stats = idx.GetTenantStats(tx, tc.tenant)
   102  					return nil
   103  				})
   104  				require.NoError(b, err)
   105  				if tc.tenant != "tenant-999" {
   106  					require.True(b, stats.DataIngested)
   107  					require.NotEqual(b, int64(0), stats.OldestProfileTime)
   108  					require.NotEqual(b, int64(0), stats.NewestProfileTime)
   109  				}
   110  			}
   111  		})
   112  	}
   113  }