github.com/grafana/pyroscope@v1.18.0/pkg/phlaredb/bucketindex/storage_test.go (about)

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  // Provenance-includes-location: https://github.com/cortexproject/cortex/blob/master/pkg/storage/tsdb/bucketindex/storage_test.go
     3  // Provenance-includes-license: Apache-2.0
     4  // Provenance-includes-copyright: The Cortex Authors.
     5  
     6  package bucketindex
     7  
     8  import (
     9  	"context"
    10  	"path"
    11  	"strings"
    12  	"testing"
    13  
    14  	"github.com/go-kit/log"
    15  	"github.com/prometheus/common/model"
    16  	"github.com/stretchr/testify/assert"
    17  	"github.com/stretchr/testify/require"
    18  
    19  	objstore_testutil "github.com/grafana/pyroscope/pkg/objstore/testutil"
    20  	"github.com/grafana/pyroscope/pkg/phlaredb/block"
    21  	block_testutil "github.com/grafana/pyroscope/pkg/phlaredb/block/testutil"
    22  )
    23  
    24  func TestReadIndex_ShouldReturnErrorIfIndexDoesNotExist(t *testing.T) {
    25  	bkt, _ := objstore_testutil.NewFilesystemBucket(t, context.Background(), t.TempDir())
    26  
    27  	idx, err := ReadIndex(context.Background(), bkt, "user-1", nil, log.NewNopLogger())
    28  	require.Equal(t, ErrIndexNotFound, err)
    29  	require.Nil(t, idx)
    30  }
    31  
    32  func TestReadIndex_ShouldReturnErrorIfIndexIsCorrupted(t *testing.T) {
    33  	const userID = "user-1"
    34  
    35  	ctx := context.Background()
    36  	bkt, _ := objstore_testutil.NewFilesystemBucket(t, ctx, t.TempDir())
    37  
    38  	// Write a corrupted index.
    39  	require.NoError(t, bkt.Upload(ctx, path.Join(userID, "phlaredb/", IndexCompressedFilename), strings.NewReader("invalid!}")))
    40  
    41  	idx, err := ReadIndex(ctx, bkt, userID, nil, log.NewNopLogger())
    42  	require.Equal(t, ErrIndexCorrupted, err)
    43  	require.Nil(t, idx)
    44  }
    45  
    46  func TestReadIndex_ShouldReturnTheParsedIndexOnSuccess(t *testing.T) {
    47  	const userID = "user-1"
    48  
    49  	ctx := context.Background()
    50  	logger := log.NewNopLogger()
    51  
    52  	bkt, _ := objstore_testutil.NewFilesystemBucket(t, ctx, t.TempDir())
    53  
    54  	// Mock some blocks in the storage.
    55  	bkt = block.BucketWithGlobalMarkers(bkt)
    56  	block_testutil.MockStorageBlock(t, bkt, userID, 10, 20)
    57  	block_testutil.MockStorageBlock(t, bkt, userID, 20, 30)
    58  	block_testutil.MockStorageDeletionMark(t, bkt, userID, block_testutil.MockStorageBlock(t, bkt, userID, 30, 40))
    59  
    60  	// Write the index.
    61  	u := NewUpdater(bkt, userID, nil, logger)
    62  	expectedIdx, _, err := u.UpdateIndex(ctx, nil)
    63  	require.NoError(t, err)
    64  	require.NoError(t, WriteIndex(ctx, bkt, userID, nil, expectedIdx))
    65  
    66  	// Read it back and compare.
    67  	actualIdx, err := ReadIndex(ctx, bkt, userID, nil, logger)
    68  	require.NoError(t, err)
    69  	assert.Equal(t, expectedIdx, actualIdx)
    70  }
    71  
    72  func BenchmarkReadIndex(b *testing.B) {
    73  	const (
    74  		numBlocks             = 1000
    75  		numBlockDeletionMarks = 100
    76  		userID                = "user-1"
    77  	)
    78  
    79  	ctx := context.Background()
    80  	logger := log.NewNopLogger()
    81  
    82  	bkt, _ := objstore_testutil.NewFilesystemBucket(b, ctx, b.TempDir())
    83  
    84  	// Mock some blocks and deletion marks in the storage.
    85  	bkt = block.BucketWithGlobalMarkers(bkt)
    86  	for i := 0; i < numBlocks; i++ {
    87  		minT := int64(i * 10)
    88  		maxT := int64((i + 1) * 10)
    89  
    90  		meta := block_testutil.MockStorageBlock(b, bkt, userID, model.Time(minT), model.Time(maxT))
    91  
    92  		if i < numBlockDeletionMarks {
    93  			block_testutil.MockStorageDeletionMark(b, bkt, userID, meta)
    94  		}
    95  	}
    96  
    97  	// Write the index.
    98  	u := NewUpdater(bkt, userID, nil, logger)
    99  	idx, _, err := u.UpdateIndex(ctx, nil)
   100  	require.NoError(b, err)
   101  	require.NoError(b, WriteIndex(ctx, bkt, userID, nil, idx))
   102  
   103  	// Read it back once just to make sure the index contains the expected data.
   104  	idx, err = ReadIndex(ctx, bkt, userID, nil, logger)
   105  	require.NoError(b, err)
   106  	require.Len(b, idx.Blocks, numBlocks)
   107  	require.Len(b, idx.BlockDeletionMarks, numBlockDeletionMarks)
   108  
   109  	b.ResetTimer()
   110  
   111  	for n := 0; n < b.N; n++ {
   112  		_, err := ReadIndex(ctx, bkt, userID, nil, logger)
   113  		require.NoError(b, err)
   114  	}
   115  }
   116  
   117  func TestDeleteIndex_ShouldNotReturnErrorIfIndexDoesNotExist(t *testing.T) {
   118  	ctx := context.Background()
   119  	bkt, _ := objstore_testutil.NewFilesystemBucket(t, ctx, t.TempDir())
   120  
   121  	assert.NoError(t, DeleteIndex(ctx, bkt, "user-1", nil))
   122  }