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

     1  package tombstones
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/hashicorp/raft"
     8  	"github.com/stretchr/testify/assert"
     9  	"github.com/stretchr/testify/require"
    10  
    11  	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
    12  	"github.com/grafana/pyroscope/pkg/metastore/index/tombstones/store"
    13  	"github.com/grafana/pyroscope/pkg/test"
    14  )
    15  
    16  func TestTombstonesRestore(t *testing.T) {
    17  	now := time.Now()
    18  	db := test.BoltDB(t)
    19  	tombstoneStore := store.NewTombstoneStore()
    20  
    21  	ts := NewTombstones(tombstoneStore, nil)
    22  	tx, err := db.Begin(true)
    23  	require.NoError(t, err)
    24  	require.NoError(t, ts.Init(tx))
    25  	require.NoError(t, tx.Commit())
    26  
    27  	for _, tombstone := range []*metastorev1.Tombstones{
    28  		{
    29  			Blocks: &metastorev1.BlockTombstones{
    30  				Name:   "x-1",
    31  				Tenant: "tenant-1",
    32  				Shard:  1,
    33  				Blocks: []string{"block-1-1", "block-1-2"},
    34  			},
    35  		},
    36  		{
    37  			Blocks: &metastorev1.BlockTombstones{
    38  				Name:   "x-2",
    39  				Tenant: "tenant-1",
    40  				Shard:  1,
    41  				Blocks: []string{"block-2-1", "block-2-2"},
    42  			},
    43  		},
    44  		{
    45  			Blocks: &metastorev1.BlockTombstones{
    46  				Name:   "x-3",
    47  				Tenant: "tenant-2",
    48  				Shard:  2,
    49  				Blocks: []string{"block-3-1", "block-3-2"},
    50  			},
    51  		},
    52  	} {
    53  		tx, err := db.Begin(true)
    54  		require.NoError(t, err)
    55  		err = ts.AddTombstones(tx, &raft.Log{Index: 1, AppendedAt: now}, tombstone)
    56  		require.NoError(t, err)
    57  		require.NoError(t, tx.Commit())
    58  	}
    59  
    60  	assert.True(t, ts.Exists("tenant-1", 1, "block-1-1"))
    61  	assert.True(t, ts.Exists("tenant-1", 1, "block-1-2"))
    62  	assert.True(t, ts.Exists("tenant-1", 1, "block-2-1"))
    63  	assert.True(t, ts.Exists("tenant-1", 1, "block-2-2"))
    64  	assert.True(t, ts.Exists("tenant-2", 2, "block-3-1"))
    65  	assert.True(t, ts.Exists("tenant-2", 2, "block-3-2"))
    66  	assert.Equal(t, 3, countTombstones(ts))
    67  
    68  	restored := NewTombstones(tombstoneStore, nil)
    69  	tx, err = db.Begin(true)
    70  	require.NoError(t, err)
    71  	err = restored.Restore(tx)
    72  	require.NoError(t, err)
    73  	require.NoError(t, tx.Commit())
    74  
    75  	assert.True(t, restored.Exists("tenant-1", 1, "block-1-1"))
    76  	assert.True(t, restored.Exists("tenant-1", 1, "block-1-2"))
    77  	assert.True(t, restored.Exists("tenant-1", 1, "block-2-1"))
    78  	assert.True(t, restored.Exists("tenant-1", 1, "block-2-2"))
    79  	assert.True(t, restored.Exists("tenant-2", 2, "block-3-1"))
    80  	assert.True(t, restored.Exists("tenant-2", 2, "block-3-2"))
    81  	assert.Equal(t, 3, countTombstones(restored))
    82  
    83  	futureTime := now.Add(time.Hour)
    84  	iter := restored.ListTombstones(futureTime)
    85  	count := 0
    86  	for iter.Next() {
    87  		count++
    88  	}
    89  	assert.Equal(t, 3, count)
    90  }