github.com/grafana/pyroscope@v1.18.0/pkg/storegateway/bucket_stores_test.go (about)

     1  package storegateway
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"os"
     7  	"path/filepath"
     8  	"strings"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/go-kit/log"
    13  	"github.com/oklog/ulid/v2"
    14  	"github.com/prometheus/client_golang/prometheus"
    15  	"github.com/prometheus/common/expfmt"
    16  	"github.com/prometheus/common/model"
    17  	"github.com/stretchr/testify/assert"
    18  	"github.com/stretchr/testify/require"
    19  
    20  	ingestv1 "github.com/grafana/pyroscope/api/gen/proto/go/ingester/v1"
    21  	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
    22  	"github.com/grafana/pyroscope/pkg/objstore/providers/filesystem"
    23  	"github.com/grafana/pyroscope/pkg/phlaredb/block"
    24  	"github.com/grafana/pyroscope/pkg/test"
    25  	"github.com/grafana/pyroscope/pkg/validation"
    26  )
    27  
    28  func TestBucketStores_BlockMetricsRegistration(t *testing.T) {
    29  	ctx := context.Background()
    30  
    31  	bucketDir := filepath.Join(t.TempDir(), "bucket")
    32  	tenantDir := filepath.Join(bucketDir, "tenant-1")
    33  	phlaredbDir := filepath.Join(tenantDir, "phlaredb")
    34  	require.NoError(t, os.MkdirAll(tenantDir, 0755))
    35  	test.Copy(t, "../phlaredb/testdata", phlaredbDir)
    36  
    37  	bucket, err := filesystem.NewBucket(bucketDir)
    38  	require.NoError(t, err)
    39  
    40  	sharding := new(mockShardingStrategy)
    41  	mockLimits := validation.MockDefaultLimits()
    42  	limits, err := validation.NewOverrides(*mockLimits, nil)
    43  	require.NoError(t, err)
    44  	logger := log.NewNopLogger()
    45  	reg := prometheus.NewRegistry()
    46  	config := BucketStoreConfig{
    47  		SyncDir:               filepath.Join(t.TempDir(), "sync-dir"),
    48  		TenantSyncConcurrency: 1,
    49  		MetaSyncConcurrency:   1,
    50  	}
    51  
    52  	stores, err := NewBucketStores(config, sharding, bucket, limits, logger, reg)
    53  	require.NoError(t, err)
    54  	require.NoError(t, stores.SyncBlocks(ctx))
    55  
    56  	userStore := stores.getStore("tenant-1")
    57  	require.NotNil(t, userStore)
    58  	require.Len(t, userStore.blockSet.blocks, 3)
    59  	r, err := userStore.blockSet.blocks[0].SelectMergeByStacktraces(ctx, &ingestv1.SelectProfilesRequest{
    60  		LabelSelector: "{}",
    61  		Type: &typesv1.ProfileType{
    62  			Name:       "process_cpu",
    63  			SampleType: "cpu",
    64  			SampleUnit: "nanoseconds",
    65  			PeriodType: "cpu",
    66  			PeriodUnit: "nanoseconds",
    67  		},
    68  		Start: 0,
    69  		End:   time.Now().UnixMilli(),
    70  	}, 0)
    71  	require.NoError(t, err)
    72  	require.NotNil(t, r)
    73  
    74  	tenantBlockMetrics := []string{`
    75  # HELP pyroscopedb_block_profile_table_accesses_total Number of times a profile table was accessed
    76  # TYPE pyroscopedb_block_profile_table_accesses_total counter
    77  pyroscopedb_block_profile_table_accesses_total{table="profiles.parquet",tenant="tenant-1"} 1
    78  `, `
    79  # HELP pyroscopedb_page_reads_total Total number of pages read while querying
    80  # TYPE pyroscopedb_page_reads_total counter
    81  pyroscopedb_page_reads_total{column="Samples.list.element.StacktraceID",table="profiles",tenant="tenant-1"} 2
    82  pyroscopedb_page_reads_total{column="Samples.list.element.Value",table="profiles",tenant="tenant-1"} 2
    83  pyroscopedb_page_reads_total{column="SeriesIndex",table="profiles",tenant="tenant-1"} 2
    84  pyroscopedb_page_reads_total{column="StacktracePartition",table="profiles",tenant="tenant-1"} 2
    85  pyroscopedb_page_reads_total{column="TimeNanos",table="profiles",tenant="tenant-1"} 2
    86  `}
    87  	assertMetricsGathered(t, reg, true, tenantBlockMetrics)
    88  
    89  	assert.NoError(t, stores.closeBucketStore("tenant-1"))
    90  	assertMetricsGathered(t, reg, false, tenantBlockMetrics)
    91  }
    92  
    93  type mockShardingStrategy struct{}
    94  
    95  func (m *mockShardingStrategy) FilterUsers(_ context.Context, userIDs []string) ([]string, error) {
    96  	return userIDs, nil
    97  }
    98  
    99  func (m *mockShardingStrategy) FilterBlocks(_ context.Context, _ string, _ map[ulid.ULID]*block.Meta, _ map[ulid.ULID]struct{}, _ block.GaugeVec) error {
   100  	return nil
   101  }
   102  
   103  func assertMetricsGathered(tb testing.TB, reg prometheus.Gatherer, gathered bool, expected []string) {
   104  	m, err := reg.Gather()
   105  	require.NoError(tb, err)
   106  	actual := make([]string, 0, len(m))
   107  	var buf bytes.Buffer
   108  	for _, metric := range m {
   109  		buf.Reset()
   110  		_, err = expfmt.MetricFamilyToText(&buf, metric)
   111  		require.NoError(tb, err)
   112  		if s := strings.TrimSpace(buf.String()); len(s) > 0 {
   113  			actual = append(actual, s)
   114  		}
   115  	}
   116  
   117  	p := expfmt.NewTextParser(model.UTF8Validation)
   118  	for _, e := range expected {
   119  		_, err = p.TextToMetricFamilies(strings.NewReader(e))
   120  		require.NoError(tb, err, "expected metric is not valid:\n%s", e)
   121  		e = strings.TrimSpace(e)
   122  		var found bool
   123  		for _, g := range actual {
   124  			if found = g == e; found {
   125  				break
   126  			}
   127  		}
   128  		if gathered {
   129  			assert.True(tb, found, "expected metric not found:\n%s", e)
   130  		} else {
   131  			assert.False(tb, found, "found unexpected metric:\n%s", e)
   132  		}
   133  	}
   134  
   135  	if tb.Failed() {
   136  		tb.Logf("gathered metrics:\n%s\n", strings.Join(actual, "\n\n"))
   137  	}
   138  }