github.com/grafana/pyroscope@v1.18.0/pkg/metastore/compaction/compactor/compactor_test.go (about)

     1  package compactor
     2  
     3  import (
     4  	"errors"
     5  	"strconv"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/hashicorp/raft"
    10  	"github.com/stretchr/testify/mock"
    11  	"github.com/stretchr/testify/require"
    12  
    13  	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
    14  	"github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1/raft_log"
    15  	"github.com/grafana/pyroscope/pkg/iter"
    16  	"github.com/grafana/pyroscope/pkg/metastore/compaction"
    17  	"github.com/grafana/pyroscope/pkg/test"
    18  	"github.com/grafana/pyroscope/pkg/test/mocks/mockcompactor"
    19  )
    20  
    21  func TestCompactor_Compact(t *testing.T) {
    22  	queueStore := new(mockcompactor.MockBlockQueueStore)
    23  	tombstones := new(mockcompactor.MockTombstones)
    24  
    25  	e := compaction.BlockEntry{
    26  		Index:      1,
    27  		AppendedAt: time.Unix(0, 0).UnixNano(),
    28  		ID:         "1",
    29  		Tenant:     "A",
    30  	}
    31  
    32  	compactor := NewCompactor(testConfig, queueStore, tombstones, nil)
    33  	testErr := errors.New("x")
    34  	t.Run("fails if cannot store the entry", test.AssertIdempotentSubtest(t, func(t *testing.T) {
    35  		queueStore.On("StoreEntry", mock.Anything, mock.Anything).Return(testErr)
    36  		require.ErrorIs(t, compactor.Compact(nil, e), testErr)
    37  	}))
    38  
    39  	queueStore.AssertExpectations(t)
    40  	tombstones.AssertExpectations(t)
    41  }
    42  
    43  func TestCompactor_UpdatePlan(t *testing.T) {
    44  	const N = 10
    45  
    46  	tombstones := new(mockcompactor.MockTombstones)
    47  	queueStore := new(mockcompactor.MockBlockQueueStore)
    48  	queueStore.On("StoreEntry", mock.Anything, mock.Anything).
    49  		Return(nil).Times(N)
    50  
    51  	compactor := NewCompactor(testConfig, queueStore, tombstones, nil)
    52  	now := time.Unix(0, 0)
    53  	for i := 0; i < N; i++ {
    54  		err := compactor.Compact(nil, compaction.BlockEntry{
    55  			Index:      1,
    56  			AppendedAt: now.UnixNano(),
    57  			ID:         strconv.Itoa(i),
    58  			Tenant:     "A",
    59  		})
    60  		require.NoError(t, err)
    61  	}
    62  
    63  	planned := make([]*raft_log.CompactionJobPlan, 3)
    64  	test.AssertIdempotent(t, func(t *testing.T) {
    65  		tombstones.On("ListTombstones", mock.Anything).
    66  			Return(iter.NewEmptyIterator[*metastorev1.Tombstones](), nil)
    67  
    68  		planner := compactor.NewPlan(&raft.Log{Index: uint64(2), AppendedAt: now})
    69  		for i := range planned {
    70  			job, err := planner.CreateJob()
    71  			require.NoError(t, err)
    72  			require.NotNil(t, job)
    73  			planned[i] = job
    74  		}
    75  
    76  		job, err := planner.CreateJob()
    77  		require.NoError(t, err)
    78  		require.Nil(t, job)
    79  	})
    80  
    81  	// UpdatePlan is mostly idempotent, except it won't
    82  	// DeleteEntry that is not loaded into memory.
    83  	queueStore.On("DeleteEntry", mock.Anything, mock.Anything, mock.Anything).
    84  		Return(nil).Times(9)
    85  
    86  	test.AssertIdempotent(t, func(t *testing.T) {
    87  		newJobs := make([]*raft_log.NewCompactionJob, 3)
    88  		for i := range planned {
    89  			newJobs[i] = &raft_log.NewCompactionJob{Plan: planned[i]}
    90  		}
    91  
    92  		update := &raft_log.CompactionPlanUpdate{NewJobs: newJobs}
    93  		require.NoError(t, compactor.UpdatePlan(nil, update))
    94  
    95  		planner := compactor.NewPlan(&raft.Log{Index: uint64(3), AppendedAt: now})
    96  		job, err := planner.CreateJob()
    97  		require.NoError(t, err)
    98  		require.Nil(t, job)
    99  	})
   100  
   101  	queueStore.AssertExpectations(t)
   102  	tombstones.AssertExpectations(t)
   103  }
   104  
   105  func TestCompactor_Restore(t *testing.T) {
   106  	queueStore := new(mockcompactor.MockBlockQueueStore)
   107  	queueStore.On("ListEntries", mock.Anything).Return(iter.NewSliceIterator([]compaction.BlockEntry{
   108  		{Index: 0, ID: "0", Tenant: "A"},
   109  		{Index: 1, ID: "1", Tenant: "A"},
   110  		{Index: 2, ID: "2", Tenant: "A"},
   111  		{Index: 3, ID: "3", Tenant: "A"},
   112  	}))
   113  
   114  	tombstones := new(mockcompactor.MockTombstones)
   115  	tombstones.On("ListTombstones", mock.Anything).
   116  		Return(iter.NewEmptyIterator[*metastorev1.Tombstones](), nil)
   117  
   118  	compactor := NewCompactor(testConfig, queueStore, tombstones, nil)
   119  	require.NoError(t, compactor.Restore(nil))
   120  
   121  	planner := compactor.NewPlan(new(raft.Log))
   122  	planned, err := planner.CreateJob()
   123  	require.NoError(t, err)
   124  	require.NotEmpty(t, planned)
   125  
   126  	queueStore.AssertExpectations(t)
   127  	tombstones.AssertExpectations(t)
   128  }