github.com/grafana/pyroscope@v1.18.0/pkg/operations/handlers_test.go (about)

     1  package operations
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"encoding/json"
     7  	"fmt"
     8  	"net/http/httptest"
     9  	"strings"
    10  	"testing"
    11  	"time"
    12  
    13  	"github.com/dustin/go-humanize"
    14  	"github.com/go-kit/log"
    15  	"github.com/gorilla/mux"
    16  	"github.com/grafana/dskit/concurrency"
    17  	"github.com/oklog/ulid/v2"
    18  	"github.com/prometheus/common/model"
    19  	"github.com/stretchr/testify/assert"
    20  	"github.com/stretchr/testify/require"
    21  
    22  	"github.com/grafana/pyroscope/pkg/objstore/testutil"
    23  	"github.com/grafana/pyroscope/pkg/phlaredb/block"
    24  	"github.com/grafana/pyroscope/pkg/phlaredb/bucketindex"
    25  )
    26  
    27  func TestHandlers_CreateBlockDetailsHandler(t *testing.T) {
    28  	const userID = "user-1"
    29  
    30  	ctx := context.Background()
    31  
    32  	bkt, _ := testutil.NewFilesystemBucket(t, ctx, t.TempDir())
    33  	now := time.Now()
    34  	logs := &concurrency.SyncBuffer{}
    35  	logger := log.NewLogfmtLogger(logs)
    36  
    37  	// Create a bucket index.
    38  	blk := &bucketindex.Block{ID: ulid.MustNew(1, nil), MinTime: model.Now().Add(-2 * time.Hour), MaxTime: model.Now()}
    39  	blkMeta := block.Meta{ULID: blk.ID, Compaction: block.BlockMetaCompaction{Level: 3}, Version: block.MetaVersion3}
    40  	metaJson, _ := json.Marshal(blkMeta)
    41  	require.NoError(t, bkt.Upload(ctx, fmt.Sprintf("user-1/phlaredb/%s/meta.json", blk.ID.String()), bytes.NewReader(metaJson)))
    42  
    43  	require.NoError(t, bucketindex.WriteIndex(ctx, bkt, userID, nil, &bucketindex.Index{
    44  		Version:            bucketindex.IndexVersion3,
    45  		Blocks:             bucketindex.Blocks{blk},
    46  		BlockDeletionMarks: bucketindex.BlockDeletionMarks{},
    47  		UpdatedAt:          now.Unix(),
    48  	}))
    49  
    50  	handlers := Handlers{
    51  		Bucket: bkt,
    52  		Logger: logger,
    53  	}
    54  
    55  	m := mux.NewRouter()
    56  	m.HandleFunc("/tenants/{tenant}/blocks/{block}", handlers.CreateBlockDetailsHandler())
    57  
    58  	req := httptest.NewRequest("GET", "/tenants/user-1/blocks/"+blk.ID.String(), nil)
    59  	resp := httptest.NewRecorder()
    60  
    61  	m.ServeHTTP(resp, req)
    62  
    63  	require.Equal(t, 200, resp.Code)
    64  	require.True(t, strings.Contains(resp.Body.String(), blk.ID.String()))
    65  	require.True(t, strings.Contains(resp.Body.String(), "<td>3</td>"))
    66  }
    67  
    68  func TestHandlers_CreateBlocksHandler(t *testing.T) {
    69  	const userID = "user-1"
    70  
    71  	ctx := context.Background()
    72  
    73  	bkt, _ := testutil.NewFilesystemBucket(t, ctx, t.TempDir())
    74  	now := time.Now()
    75  	logs := &concurrency.SyncBuffer{}
    76  	logger := log.NewLogfmtLogger(logs)
    77  
    78  	// Create a bucket index.
    79  	block1 := &bucketindex.Block{ID: ulid.MustNew(1, nil), MinTime: model.Now().Add(-2 * time.Hour), MaxTime: model.Now()}
    80  	block2 := &bucketindex.Block{ID: ulid.MustNew(2, nil), MinTime: model.Now().Add(-4 * time.Hour), MaxTime: model.Now()}
    81  	block3 := &bucketindex.Block{ID: ulid.MustNew(3, nil), MinTime: model.Now().Add(-4 * time.Hour), MaxTime: model.Now()}
    82  	block4 := &bucketindex.Block{ID: ulid.MustNew(4, nil), MinTime: model.Now().Add(-12 * time.Hour), MaxTime: model.Now().Add(-10 * time.Hour)}
    83  
    84  	require.NoError(t, bucketindex.WriteIndex(ctx, bkt, userID, nil, &bucketindex.Index{
    85  		Version:            bucketindex.IndexVersion1,
    86  		Blocks:             bucketindex.Blocks{block1, block2, block3, block4},
    87  		BlockDeletionMarks: bucketindex.BlockDeletionMarks{},
    88  		UpdatedAt:          now.Unix(),
    89  	}))
    90  
    91  	handlers := Handlers{
    92  		Bucket: bkt,
    93  		Logger: logger,
    94  	}
    95  
    96  	m := mux.NewRouter()
    97  	m.HandleFunc("/tenants/{tenant}/blocks", handlers.CreateBlocksHandler())
    98  
    99  	req := httptest.NewRequest("GET", "/tenants/user-1/blocks?queryFrom=now-4h", nil)
   100  	resp := httptest.NewRecorder()
   101  
   102  	m.ServeHTTP(resp, req)
   103  
   104  	require.Equal(t, 200, resp.Code)
   105  	require.True(t, strings.Contains(resp.Body.String(), block1.ID.String()))
   106  	require.True(t, strings.Contains(resp.Body.String(), block2.ID.String()))
   107  	require.True(t, strings.Contains(resp.Body.String(), block3.ID.String()))
   108  	require.False(t, strings.Contains(resp.Body.String(), block4.ID.String())) // outside the now-4h window
   109  }
   110  
   111  func TestHandlers_CreateIndexHandler(t *testing.T) {
   112  	const userID = "user-1"
   113  
   114  	ctx := context.Background()
   115  
   116  	bkt, _ := testutil.NewFilesystemBucket(t, ctx, t.TempDir())
   117  	logs := &concurrency.SyncBuffer{}
   118  	logger := log.NewLogfmtLogger(logs)
   119  
   120  	// Create a bucket index.
   121  	require.NoError(t, bucketindex.WriteIndex(ctx, bkt, userID, nil, &bucketindex.Index{
   122  		Version:            bucketindex.IndexVersion3,
   123  		Blocks:             bucketindex.Blocks{},
   124  		BlockDeletionMarks: bucketindex.BlockDeletionMarks{},
   125  	}))
   126  
   127  	handlers := Handlers{
   128  		Bucket: bkt,
   129  		Logger: logger,
   130  	}
   131  
   132  	h := handlers.CreateIndexHandler()
   133  
   134  	req := httptest.NewRequest("GET", "/", nil)
   135  	resp := httptest.NewRecorder()
   136  
   137  	h(resp, req)
   138  
   139  	require.Equal(t, 200, resp.Code)
   140  	require.True(t, strings.Contains(resp.Body.String(), "user-1"))
   141  }
   142  
   143  func Test_filterAndGroupBlocks(t *testing.T) {
   144  	now := model.TimeFromUnixNano(time.Date(2025, 10, 16, 16, 0, 0, 0, time.UTC).UnixNano())
   145  	block1 := &bucketindex.Block{ID: ulid.MustNew(1, nil), MinTime: now.Add(-2 * time.Hour), MaxTime: now.Add(-1 * time.Hour)}
   146  	block2 := &bucketindex.Block{ID: ulid.MustNew(2, nil), MinTime: now.Add(-4 * time.Hour), MaxTime: now.Add(-3 * time.Hour)}
   147  	block3 := &bucketindex.Block{ID: ulid.MustNew(3, nil), MinTime: now.Add(-4*time.Hour + time.Minute), MaxTime: now.Add(-3 * time.Hour)}
   148  	block4 := &bucketindex.Block{ID: ulid.MustNew(4, nil), MinTime: now.Add(-12 * time.Hour), MaxTime: now.Add(-10 * time.Hour)}
   149  	h := &Handlers{MaxBlockDuration: time.Hour}
   150  
   151  	type args struct {
   152  		index *bucketindex.Index
   153  		query *blockQuery
   154  	}
   155  	tests := []struct {
   156  		name string
   157  		args args
   158  		want *blockListResult
   159  	}{
   160  		{
   161  			name: "empty index",
   162  			args: args{
   163  				index: &bucketindex.Index{
   164  					Version:            bucketindex.IndexVersion3,
   165  					Blocks:             bucketindex.Blocks{},
   166  					BlockDeletionMarks: bucketindex.BlockDeletionMarks{},
   167  				}, query: &blockQuery{}},
   168  			want: &blockListResult{
   169  				BlockGroups:          []*blockGroup{},
   170  				GroupDurationMinutes: 0,
   171  			},
   172  		},
   173  		{
   174  			name: "index with blocks that can be filtered and grouped",
   175  			args: args{
   176  				index: &bucketindex.Index{
   177  					Version:            bucketindex.IndexVersion3,
   178  					Blocks:             bucketindex.Blocks{block1, block2, block3, block4},
   179  					BlockDeletionMarks: bucketindex.BlockDeletionMarks{&bucketindex.BlockDeletionMark{ID: block1.ID}},
   180  				},
   181  				query: &blockQuery{
   182  					parsedFrom: now.Time().Add(-6 * time.Hour),
   183  					parsedTo:   now.Time(),
   184  				}},
   185  			want: &blockListResult{
   186  				// block 1 is not included because it is marked as deleted
   187  				// block 4 is not included because it is outside the query window
   188  				BlockGroups: []*blockGroup{
   189  					{
   190  						MinTime:          block2.MinTime.Time().Truncate(time.Hour).UTC(),
   191  						FormattedMinTime: block2.MinTime.Time().Truncate(time.Hour).UTC().Format(time.RFC3339),
   192  						Blocks: []*blockDetails{
   193  							{
   194  								ID:                block3.ID.String(), // block 3 is newer so it goes first
   195  								MinTime:           block3.MinTime.Time().UTC().Format(time.RFC3339),
   196  								MaxTime:           block3.MaxTime.Time().UTC().Format(time.RFC3339),
   197  								Duration:          59,
   198  								FormattedDuration: "59m0s",
   199  								UploadedAt:        time.UnixMilli(0).UTC().Format(time.RFC3339),
   200  							},
   201  							{
   202  								ID:                block2.ID.String(),
   203  								MinTime:           block2.MinTime.Time().UTC().Format(time.RFC3339),
   204  								MaxTime:           block2.MaxTime.Time().UTC().Format(time.RFC3339),
   205  								Duration:          60,
   206  								FormattedDuration: "1h0m0s",
   207  								UploadedAt:        time.UnixMilli(0).UTC().Format(time.RFC3339),
   208  							},
   209  						},
   210  						MinTimeAge:              humanize.RelTime(block2.MinTime.Time(), now.Time(), "ago", ""),
   211  						MaxBlockDurationMinutes: 60,
   212  					},
   213  				},
   214  				GroupDurationMinutes: 60,
   215  				MaxBlocksPerGroup:    2,
   216  			},
   217  		},
   218  	}
   219  	for _, tt := range tests {
   220  		t.Run(tt.name, func(t *testing.T) {
   221  			assert.Equalf(t, tt.want, h.filterAndGroupBlocks(tt.args.index, tt.args.query, now.Time()), "filterAndGroupBlocks(%v, %v)", tt.args.index, tt.args.query)
   222  		})
   223  	}
   224  }