github.com/grafana/pyroscope@v1.18.0/pkg/phlaredb/block/testutil/create_block.go (about)

     1  package testutil
     2  
     3  import (
     4  	"context"
     5  	"crypto/rand"
     6  	"encoding/json"
     7  	"os"
     8  	"path/filepath"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/oklog/ulid/v2"
    13  	"github.com/stretchr/testify/assert"
    14  
    15  	"github.com/grafana/pyroscope/pkg/objstore/providers/filesystem"
    16  
    17  	"github.com/prometheus/common/model"
    18  	"github.com/stretchr/testify/require"
    19  
    20  	phlaremodel "github.com/grafana/pyroscope/pkg/model"
    21  	"github.com/grafana/pyroscope/pkg/phlaredb"
    22  	"github.com/grafana/pyroscope/pkg/phlaredb/block"
    23  	"github.com/grafana/pyroscope/pkg/pprof/testhelper"
    24  )
    25  
    26  type noLimit struct{}
    27  
    28  func (n noLimit) AllowProfile(fp model.Fingerprint, lbs phlaremodel.Labels, tsNano int64) error {
    29  	return nil
    30  }
    31  
    32  func (n noLimit) Stop() {}
    33  
    34  // CreateBlock creates a block with the given profiles.
    35  // Returns the block metadata, the directory where the block is stored, and an error if any.
    36  func CreateBlock(t testing.TB, generator func() []*testhelper.ProfileBuilder) (block.Meta, string) {
    37  	t.Helper()
    38  	dir := t.TempDir()
    39  	ctx := context.Background()
    40  	h, err := phlaredb.NewHead(ctx, phlaredb.Config{
    41  		DataPath:         dir,
    42  		MaxBlockDuration: 24 * time.Hour,
    43  		Parquet: &phlaredb.ParquetConfig{
    44  			MaxBufferRowCount: 10,
    45  		},
    46  	}, noLimit{})
    47  	require.NoError(t, err)
    48  
    49  	// ingest.
    50  	for _, p := range generator() {
    51  		require.NoError(t, h.Ingest(ctx, p.Profile, p.UUID, nil, p.Labels...))
    52  	}
    53  
    54  	require.NoError(t, h.Flush(ctx))
    55  	require.NoError(t, h.Move())
    56  	localDir := filepath.Join(dir, phlaredb.PathLocal)
    57  	metaMap, err := block.ListBlocks(localDir, time.Time{})
    58  	require.NoError(t, err)
    59  	require.Len(t, metaMap, 1)
    60  	var meta *block.Meta
    61  	for _, m := range metaMap {
    62  		meta = m
    63  	}
    64  	require.NotNil(t, meta)
    65  	return *meta, localDir
    66  }
    67  
    68  func OpenBlockFromMemory(t *testing.T, dir string, minTime, maxTime model.Time, profiles, tsdb, symbols []byte) *phlaredb.BlockQuerier {
    69  	CreateBlockFromMemory(t, dir, minTime, maxTime, profiles, tsdb, symbols)
    70  	blockBucket, err := filesystem.NewBucket(dir)
    71  	require.NoError(t, err)
    72  	blockQuerier := phlaredb.NewBlockQuerier(context.Background(), blockBucket)
    73  
    74  	err = blockQuerier.Sync(context.Background())
    75  	require.NoError(t, err)
    76  
    77  	return blockQuerier
    78  }
    79  
    80  func CreateBlockFromMemory(t *testing.T, dir string, minTime, maxTime model.Time, profiles, tsdb, symbols []byte) *block.Meta {
    81  	blockid, err := ulid.New(uint64(maxTime), rand.Reader)
    82  	require.NoError(t, err)
    83  	blockDir := filepath.Join(dir, blockid.String())
    84  	err = os.MkdirAll(blockDir, 0755)
    85  	assert.NoError(t, err)
    86  	err = os.WriteFile(filepath.Join(blockDir, "profiles.parquet"), profiles, 0644)
    87  	assert.NoError(t, err)
    88  	err = os.WriteFile(filepath.Join(blockDir, "index.tsdb"), tsdb, 0644)
    89  	assert.NoError(t, err)
    90  	err = os.WriteFile(filepath.Join(blockDir, "symbols.symdb"), symbols, 0644)
    91  	assert.NoError(t, err)
    92  
    93  	blockMeta := &block.Meta{
    94  		ULID:    blockid,
    95  		MinTime: minTime,
    96  		MaxTime: maxTime,
    97  		Files: []block.File{
    98  			{
    99  				RelPath:   "profiles.parquet",
   100  				SizeBytes: uint64(len(profiles)),
   101  			},
   102  			{
   103  				RelPath:   "index.tsdb",
   104  				SizeBytes: uint64(len(tsdb)),
   105  			},
   106  			{
   107  				RelPath:   "symbols.symdb",
   108  				SizeBytes: uint64(len(symbols)),
   109  			},
   110  		},
   111  		Version: block.MetaVersion3,
   112  	}
   113  	blockMetaJson, err := json.Marshal(&blockMeta)
   114  	assert.NoError(t, err)
   115  	err = os.WriteFile(filepath.Join(blockDir, block.MetaFilename), blockMetaJson, 0644)
   116  	assert.NoError(t, err)
   117  
   118  	return blockMeta
   119  }