github.com/grafana/pyroscope@v1.18.0/pkg/metastore/fsm/boltdb_test.go (about)

     1  package fsm
     2  
     3  import (
     4  	"bytes"
     5  	"os"
     6  	"path/filepath"
     7  	"strings"
     8  	"testing"
     9  
    10  	"github.com/go-kit/log"
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/require"
    13  	"go.etcd.io/bbolt"
    14  )
    15  
    16  func TestBoltDB_open_restore(t *testing.T) {
    17  	tempDir := t.TempDir()
    18  	buf := bytes.NewBuffer(nil)
    19  	db := newDB(log.NewLogfmtLogger(buf), newMetrics(nil), Config{DataDir: tempDir})
    20  	require.NoError(t, db.open(false))
    21  
    22  	data := []string{
    23  		"k1", "v1",
    24  		"k2", "v2",
    25  		"k3", "v3",
    26  	}
    27  
    28  	snapshotSource := filepath.Join(tempDir, "snapshot_source")
    29  	require.NoError(t, createDB(t, snapshotSource, data).Close())
    30  	s, err := os.ReadFile(snapshotSource)
    31  	require.NoError(t, err)
    32  	require.NoError(t, db.restore(bytes.NewReader(s)))
    33  
    34  	collected := make([]string, 0, len(data))
    35  	require.NoError(t, db.boltdb.View(func(tx *bbolt.Tx) error {
    36  		b := tx.Bucket([]byte("test"))
    37  		assert.NotNil(t, b)
    38  		return b.ForEach(func(k, v []byte) error {
    39  			collected = append(collected, string(k), string(v))
    40  			return nil
    41  		})
    42  	}))
    43  
    44  	assert.Equal(t, data, collected)
    45  	assert.False(t, strings.Contains(buf.String(), "compacting snapshot"))
    46  }
    47  
    48  func TestBoltDB_open_restore_compact(t *testing.T) {
    49  	tempDir := t.TempDir()
    50  	buf := bytes.NewBuffer(nil)
    51  	db := newDB(log.NewLogfmtLogger(buf), newMetrics(nil), Config{
    52  		DataDir:                  tempDir,
    53  		SnapshotCompactOnRestore: true,
    54  	})
    55  	require.NoError(t, db.open(false))
    56  
    57  	data := []string{
    58  		"k1", "v1",
    59  		"k2", "v2",
    60  		"k3", "v3",
    61  	}
    62  
    63  	snapshotSource := filepath.Join(tempDir, "snapshot_source")
    64  	require.NoError(t, createDB(t, snapshotSource, data).Close())
    65  	s, err := os.ReadFile(snapshotSource)
    66  	require.NoError(t, err)
    67  	require.NoError(t, db.restore(bytes.NewReader(s)))
    68  
    69  	collected := make([]string, 0, len(data))
    70  	require.NoError(t, db.boltdb.View(func(tx *bbolt.Tx) error {
    71  		b := tx.Bucket([]byte("test"))
    72  		assert.NotNil(t, b)
    73  		return b.ForEach(func(k, v []byte) error {
    74  			collected = append(collected, string(k), string(v))
    75  			return nil
    76  		})
    77  	}))
    78  
    79  	assert.Equal(t, data, collected)
    80  	assert.True(t, strings.Contains(buf.String(), "compacting snapshot"))
    81  }
    82  
    83  func createDB(t *testing.T, path string, pairs []string) *bbolt.DB {
    84  	opts := bbolt.Options{
    85  		NoGrowSync:     true,
    86  		NoFreelistSync: true,
    87  		NoSync:         true,
    88  		FreelistType:   bbolt.FreelistMapType,
    89  	}
    90  	db, err := bbolt.Open(path, 0644, &opts)
    91  	require.NoError(t, err)
    92  	require.NoError(t, db.Update(func(tx *bbolt.Tx) error {
    93  		bucket, err := tx.CreateBucketIfNotExists([]byte("test"))
    94  		require.NoError(t, err)
    95  		for len(pairs) > 1 {
    96  			if err = bucket.Put([]byte(pairs[0]), []byte(pairs[1])); err != nil {
    97  				return err
    98  			}
    99  			pairs = pairs[2:]
   100  		}
   101  		return nil
   102  	}))
   103  	return db
   104  }