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

     1  package dlq
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/prometheus/client_golang/prometheus"
    10  	"github.com/prometheus/client_golang/prometheus/testutil"
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/mock"
    13  	"github.com/stretchr/testify/require"
    14  	"google.golang.org/grpc/codes"
    15  	"google.golang.org/grpc/status"
    16  
    17  	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
    18  	"github.com/grafana/pyroscope/pkg/block"
    19  	"github.com/grafana/pyroscope/pkg/metastore/raftnode/raftnodepb"
    20  	"github.com/grafana/pyroscope/pkg/objstore/providers/memory"
    21  	"github.com/grafana/pyroscope/pkg/test"
    22  	"github.com/grafana/pyroscope/pkg/test/mocks/mockdlq"
    23  )
    24  
    25  func TestRecoverTick(t *testing.T) {
    26  	metas := []*metastorev1.BlockMeta{
    27  		{
    28  			Id:    test.ULID("2024-09-23T03:00:00Z"),
    29  			Shard: 2,
    30  		},
    31  		{
    32  			Id:    test.ULID("2024-09-23T01:00:00Z"),
    33  			Shard: 1,
    34  		},
    35  		{
    36  			Id:    test.ULID("2024-09-23T02:00:00Z"),
    37  			Shard: 2,
    38  		},
    39  	}
    40  
    41  	var actual []*metastorev1.BlockMeta
    42  	srv := mockdlq.NewMockMetastore(t)
    43  	srv.On("AddRecoveredBlock", mock.Anything, mock.Anything).
    44  		Times(3).
    45  		Run(func(args mock.Arguments) {
    46  			meta := args.Get(1).(*metastorev1.AddBlockRequest).Block
    47  			actual = append(actual, meta)
    48  		}).
    49  		Return(&metastorev1.AddBlockResponse{}, nil)
    50  
    51  	bucket := memory.NewInMemBucket()
    52  	for _, meta := range metas {
    53  		addMeta(bucket, meta)
    54  	}
    55  
    56  	r := NewRecovery(test.NewTestingLogger(t), Config{}, srv, bucket, prometheus.NewRegistry())
    57  	r.recoverTick(context.Background())
    58  
    59  	expected := []*metastorev1.BlockMeta{
    60  		metas[1],
    61  		metas[2],
    62  		metas[0],
    63  	}
    64  
    65  	require.Equal(t, len(actual), len(expected))
    66  	for i := range actual {
    67  		require.Equal(t, actual[i].Id, expected[i].Id)
    68  		require.Equal(t, actual[i].Shard, expected[i].Shard)
    69  	}
    70  
    71  	assert.Equal(t, 3.0, testutil.ToFloat64(r.metrics.recoveryAttempts.WithLabelValues("success")))
    72  	assert.Equal(t, 0.0, testutil.ToFloat64(r.metrics.recoveryAttempts.WithLabelValues("unmarshal_error")))
    73  	assert.Equal(t, 0.0, testutil.ToFloat64(r.metrics.recoveryAttempts.WithLabelValues("invalid_metadata")))
    74  }
    75  
    76  func TestNotRaftLeader(t *testing.T) {
    77  	metas := []*metastorev1.BlockMeta{
    78  		{
    79  			Id:    test.ULID("2024-09-23T01:00:00Z"),
    80  			Shard: 2,
    81  		},
    82  	}
    83  
    84  	srv := mockdlq.NewMockMetastore(t)
    85  	s, _ := status.New(codes.Unavailable, "mock metastore error").WithDetails(&raftnodepb.RaftNode{
    86  		Id:      "foo",
    87  		Address: "bar",
    88  	})
    89  	srv.On("AddRecoveredBlock", mock.Anything, mock.Anything).
    90  		Once().
    91  		Return(nil, s.Err())
    92  
    93  	bucket := memory.NewInMemBucket()
    94  	for _, meta := range metas {
    95  		addMeta(bucket, meta)
    96  	}
    97  
    98  	r := NewRecovery(test.NewTestingLogger(t), Config{}, srv, bucket, prometheus.NewRegistry())
    99  	r.recoverTick(context.Background())
   100  
   101  	assert.Equal(t, 1, len(bucket.Objects()))
   102  
   103  	assert.Equal(t, 1.0, testutil.ToFloat64(r.metrics.recoveryAttempts.WithLabelValues("metastore_error")))
   104  	assert.Equal(t, 0.0, testutil.ToFloat64(r.metrics.recoveryAttempts.WithLabelValues("success")))
   105  }
   106  
   107  func TestStartStop(t *testing.T) {
   108  	metas := []*metastorev1.BlockMeta{
   109  		{
   110  			Id:    test.ULID("2024-09-23T03:00:00Z"),
   111  			Shard: 2,
   112  		},
   113  		{
   114  			Id:    test.ULID("2024-09-23T01:00:00Z"),
   115  			Shard: 1,
   116  		},
   117  		{
   118  			Id:    test.ULID("2024-09-23T02:00:00Z"),
   119  			Shard: 2,
   120  		},
   121  	}
   122  	m := new(sync.Mutex)
   123  
   124  	var actual []*metastorev1.BlockMeta
   125  	srv := mockdlq.NewMockMetastore(t)
   126  	srv.On("AddRecoveredBlock", mock.Anything, mock.Anything).
   127  		Times(3).
   128  		Run(func(args mock.Arguments) {
   129  			meta := args.Get(1).(*metastorev1.AddBlockRequest).Block
   130  			m.Lock()
   131  			actual = append(actual, meta)
   132  			m.Unlock()
   133  		}).
   134  		Return(&metastorev1.AddBlockResponse{}, nil)
   135  
   136  	bucket := memory.NewInMemBucket()
   137  	for _, meta := range metas {
   138  		addMeta(bucket, meta)
   139  	}
   140  
   141  	r := NewRecovery(test.NewTestingLogger(t), Config{CheckInterval: time.Millisecond * 10}, srv, bucket, prometheus.NewRegistry())
   142  	r.Start()
   143  	defer r.Stop()
   144  
   145  	require.Eventually(t, func() bool {
   146  		m.Lock()
   147  		defer m.Unlock()
   148  		return len(actual) == 3
   149  	}, time.Second, time.Millisecond*100)
   150  
   151  	expected := []*metastorev1.BlockMeta{
   152  		metas[1],
   153  		metas[2],
   154  		metas[0],
   155  	}
   156  
   157  	require.Equal(t, len(actual), len(expected))
   158  	for i := range actual {
   159  		require.Equal(t, actual[i].Id, expected[i].Id)
   160  		require.Equal(t, actual[i].Shard, expected[i].Shard)
   161  	}
   162  
   163  	assert.Equal(t, 3.0, testutil.ToFloat64(r.metrics.recoveryAttempts.WithLabelValues("success")))
   164  }
   165  
   166  func addMeta(bucket *memory.InMemBucket, meta *metastorev1.BlockMeta) {
   167  	data, _ := meta.MarshalVT()
   168  	bucket.Set(block.MetadataDLQObjectPath(meta), data)
   169  }