feat: use pebble instand of ethdb/pebble

portalnetwork/storage/ethpepple/maxheap.go

@@ -1,39 +0,0 @@
-package ethpepple
-
-import (
-	"bytes"
-)
-
-const maxItem = 250_000 // every item has 40 bytes, so the heap most have 10MB
-
-type Item struct {
-	Distance  []byte
-	ValueSize uint64
-}
-
-type MaxHeap []Item
-
-func (m MaxHeap) Len() int {
-	return len(m)
-}
-
-func (m MaxHeap) Less(i, j int) bool {
-	// Compare Distance as byte slices
-	return bytes.Compare(m[i].Distance, m[j].Distance) > 0
-}
-
-func (m MaxHeap) Swap(i, j int) {
-	m[i], m[j] = m[j], m[i]
-}
-
-func (m *MaxHeap) Pop() interface{} {
-	old := *m
-	n := len(old)
-	item := old[n-1]
-	*m = old[0 : n-1]
-	return item
-}
-
-func (m *MaxHeap) Push(x interface{}) {
-	*m = append(*m, x.(Item))
-}

portalnetwork/storage/ethpepple/maxheap_test.go

@@ -1,33 +0,0 @@
-package ethpepple
-
-import (
-	"container/heap"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-)
-
-func TestMaxHeap(t *testing.T) {
-	expectValueSize := []uint64{40, 30, 20, 10}
-	// Create a heap and initialize it with some items
-	h := &MaxHeap{
-		{Distance: []byte{1}, ValueSize: 10},
-		{Distance: []byte{2}, ValueSize: 20},
-		{Distance: []byte{3}, ValueSize: 30},
-	}
-	heap.Init(h)
-
-	// Push a new item into the heap
-	heap.Push(h, Item{Distance: []byte{4}, ValueSize: 40})
-	heap.Push(h, Item{Distance: []byte{5}, ValueSize: 50})
-
-	removed := heap.Remove(h, 0)
-	assert.Equal(t, removed.(Item).ValueSize, uint64(50))
-
-	len := h.Len()
-	// Pop and print the largest element
-	for i := 0; i < len; i++ {
-		item := heap.Pop(h).(Item)
-		assert.Equal(t, item.ValueSize, expectValueSize[i])
-	}
-}

portalnetwork/storage/ethpepple/storage.go

@@ -1,34 +1,119 @@
 package ethpepple
 
 import (
-	"bytes"
-	"container/heap"
 	"encoding/binary"
+	"runtime"
 	"sync"
 	"sync/atomic"
 	"time"
 
-	"github.com/ethereum/go-ethereum/ethdb"
-	"github.com/ethereum/go-ethereum/ethdb/pebble"
+	"github.com/cockroachdb/pebble"
+	"github.com/cockroachdb/pebble/bloom"
+	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/p2p/enode"
 	"github.com/ethereum/go-ethereum/portalnetwork/storage"
 	"github.com/holiman/uint256"
 )
 
-const contentDeletionFraction = 0.05 // 5% of the content will be deleted when the storage capacity is hit and radius gets adjusted.
+const (
+	// minCache is the minimum amount of memory in megabytes to allocate to pebble
+	// read and write caching, split half and half.
+	minCache = 16
+
+	// minHandles is the minimum number of files handles to allocate to the open
+	// database files.
+	minHandles = 16
+
+	// 5% of the content will be deleted when the storage capacity is hit and radius gets adjusted.
+	contentDeletionFraction = 0.05
+)
 
 var _ storage.ContentStorage = &ContentStorage{}
 
 type PeppleStorageConfig struct {
 	StorageCapacityMB uint64
-	DB                ethdb.KeyValueStore
+	DB                *pebble.DB
 	NodeId            enode.ID
 	NetworkName       string
 }
 
-func NewPeppleDB(dataDir string, cache, handles int, namespace string) (ethdb.KeyValueStore, error) {
-	db, err := pebble.New(dataDir+"/"+namespace, cache, handles, namespace, false)
+func NewPeppleDB(dataDir string, cache, handles int, namespace string) (*pebble.DB, error) {
+	// Ensure we have some minimal caching and file guarantees
+	if cache < minCache {
+		cache = minCache
+	}
+	if handles < minHandles {
+		handles = minHandles
+	}
+	logger := log.New("database", namespace)
+	logger.Info("Allocated cache and file handles", "cache", common.StorageSize(cache*1024*1024), "handles", handles)
+
+	// The max memtable size is limited by the uint32 offsets stored in
+	// internal/arenaskl.node, DeferredBatchOp, and flushableBatchEntry.
+	//
+	// - MaxUint32 on 64-bit platforms;
+	// - MaxInt on 32-bit platforms.
+	//
+	// It is used when slices are limited to Uint32 on 64-bit platforms (the
+	// length limit for slices is naturally MaxInt on 32-bit platforms).
+	//
+	// Taken from https://github.com/cockroachdb/pebble/blob/master/internal/constants/constants.go
+	maxMemTableSize := (1<<31)<<(^uint(0)>>63) - 1
+
+	// Two memory tables is configured which is identical to leveldb,
+	// including a frozen memory table and another live one.
+	memTableLimit := 2
+	memTableSize := cache * 1024 * 1024 / 2 / memTableLimit
+
+	// The memory table size is currently capped at maxMemTableSize-1 due to a
+	// known bug in the pebble where maxMemTableSize is not recognized as a
+	// valid size.
+	//
+	// TODO use the maxMemTableSize as the maximum table size once the issue
+	// in pebble is fixed.
+	if memTableSize >= maxMemTableSize {
+		memTableSize = maxMemTableSize - 1
+	}
+	opt := &pebble.Options{
+		// Pebble has a single combined cache area and the write
+		// buffers are taken from this too. Assign all available
+		// memory allowance for cache.
+		Cache:        pebble.NewCache(int64(cache * 1024 * 1024)),
+		MaxOpenFiles: handles,
+
+		// The size of memory table(as well as the write buffer).
+		// Note, there may have more than two memory tables in the system.
+		MemTableSize: uint64(memTableSize),
+
+		// MemTableStopWritesThreshold places a hard limit on the size
+		// of the existent MemTables(including the frozen one).
+		// Note, this must be the number of tables not the size of all memtables
+		// according to https://github.com/cockroachdb/pebble/blob/master/options.go#L738-L742
+		// and to https://github.com/cockroachdb/pebble/blob/master/db.go#L1892-L1903.
+		MemTableStopWritesThreshold: memTableLimit,
+
+		// The default compaction concurrency(1 thread),
+		// Here use all available CPUs for faster compaction.
+		MaxConcurrentCompactions: runtime.NumCPU,
+
+		// Per-level options. Options for at least one level must be specified. The
+		// options for the last level are used for all subsequent levels.
+		Levels: []pebble.LevelOptions{
+			{TargetFileSize: 2 * 1024 * 1024, FilterPolicy: bloom.FilterPolicy(10)},
+			{TargetFileSize: 4 * 1024 * 1024, FilterPolicy: bloom.FilterPolicy(10)},
+			{TargetFileSize: 8 * 1024 * 1024, FilterPolicy: bloom.FilterPolicy(10)},
+			{TargetFileSize: 16 * 1024 * 1024, FilterPolicy: bloom.FilterPolicy(10)},
+			{TargetFileSize: 32 * 1024 * 1024, FilterPolicy: bloom.FilterPolicy(10)},
+			{TargetFileSize: 64 * 1024 * 1024, FilterPolicy: bloom.FilterPolicy(10)},
+			{TargetFileSize: 128 * 1024 * 1024, FilterPolicy: bloom.FilterPolicy(10)},
+		},
+		ReadOnly: false,
+	}
+	// Disable seek compaction explicitly. Check https://github.com/ethereum/go-ethereum/pull/20130
+	// for more details.
+	opt.Experimental.ReadSamplingMultiplier = -1
+	db, err := pebble.Open(dataDir+"/"+namespace, opt)
 	return db, err
 }
 
@@ -37,12 +122,13 @@ type ContentStorage struct {
 	storageCapacityInBytes uint64
 	radius                 atomic.Value
 	log                    log.Logger
-	db                     ethdb.KeyValueStore
+	db                     *pebble.DB
 	size                   uint64
 	sizeChan               chan uint64
 	sizeMutex              sync.RWMutex
 	isPruning              bool
 	pruneDoneChan          chan uint64 // finish prune and get the pruned size
+	writeOptions           *pebble.WriteOptions
 }
 
 func NewPeppleStorage(config PeppleStorageConfig) (storage.ContentStorage, error) {
@@ -53,17 +139,14 @@ func NewPeppleStorage(config PeppleStorageConfig) (storage.ContentStorage, error
 		log:                    log.New("storage", config.NetworkName),
 		sizeChan:               make(chan uint64, 100),
 		pruneDoneChan:          make(chan uint64, 1),
+		writeOptions:           &pebble.WriteOptions{Sync: false},
 	}
 	cs.radius.Store(storage.MaxDistance)
-	exist, err := cs.db.Has(storage.RadisuKey)
-	if err != nil {
+	radius, _, err := cs.db.Get(storage.RadisuKey)
+	if err != nil && err != pebble.ErrNotFound {
 		return nil, err
 	}
-	if exist {
-		radius, err := cs.db.Get(storage.RadisuKey)
-		if err != nil {
-			return nil, err
-		}
+	if err == nil {
 		dis := uint256.NewInt(0)
 		err = dis.UnmarshalSSZ(radius)
 		if err != nil {
@@ -72,15 +155,11 @@ func NewPeppleStorage(config PeppleStorageConfig) (storage.ContentStorage, error
 		cs.radius.Store(dis)
 	}
 
-	exist, err = cs.db.Has(storage.SizeKey)
-	if err != nil {
+	val, _, err := cs.db.Get(storage.SizeKey)
+	if err != nil && err != pebble.ErrNotFound {
 		return nil, err
 	}
-	if exist {
-		val, err := cs.db.Get(storage.SizeKey)
-		if err != nil {
-			return nil, err
-		}
+	if err == nil {
 		size := binary.BigEndian.Uint64(val)
 		// init stage, no need to use lock
 		cs.size = size
@@ -91,14 +170,24 @@ func NewPeppleStorage(config PeppleStorageConfig) (storage.ContentStorage, error
 
 // Get implements storage.ContentStorage.
 func (c *ContentStorage) Get(contentKey []byte, contentId []byte) ([]byte, error) {
-	return c.db.Get(contentId)
+	distance := xor(contentId, c.nodeId[:])
+	data, closer, err := c.db.Get(distance)
+	if err != nil && err != pebble.ErrNotFound {
+		return nil, err
+	}
+	if err == pebble.ErrNotFound {
+		return nil, storage.ErrContentNotFound
+	}
+	closer.Close()
+	return data, nil
 }
 
 // Put implements storage.ContentStorage.
 func (c *ContentStorage) Put(contentKey []byte, contentId []byte, content []byte) error {
 	length := uint64(len(contentId)) + uint64(len(content))
 	c.sizeChan <- length
-	return c.db.Put(contentId, content)
+	distance := xor(contentId, c.nodeId[:])
+	return c.db.Set(distance, content, c.writeOptions)
 }
 
 // Radius implements storage.ContentStorage.
@@ -119,7 +208,10 @@ func (c *ContentStorage) saveCapacity() {
 		case <-ticker.C:
 			if sizeChanged {
 				binary.BigEndian.PutUint64(buf, c.size)
-				c.db.Put(storage.SizeKey, buf)
+				err := c.db.Set(storage.SizeKey, buf, c.writeOptions)
+				if err != nil {
+					c.log.Error("save capacity failed", "error", err)
+				}
 				sizeChanged = false
 			}
 		case size := <-c.sizeChan:
@@ -143,65 +235,40 @@ func (c *ContentStorage) saveCapacity() {
 }
 
 func (c *ContentStorage) prune() {
-	var distance = []byte{}
-
-	h := &MaxHeap{}
-	heap.Init(h)
-
 	expectSize := uint64(float64(c.storageCapacityInBytes) * contentDeletionFraction)
-
 	var curentSize uint64 = 0
 
 	defer func() {
 		c.pruneDoneChan <- curentSize
 	}()
 	// get the keys to be deleted order by distance desc
-	iterator := c.db.NewIterator(nil, nil)
-	defer iterator.Release()
-	for iterator.Next() {
-		key := iterator.Key()
-		if bytes.Equal(key, storage.SizeKey) || bytes.Equal(key, storage.RadisuKey) {
-			continue
-		}
-		val := iterator.Value()
-		size := uint64(len(val))
-
-		distance := xor(key, c.nodeId[:])
-		heap.Push(h, Item{
-			Distance:  distance,
-			ValueSize: size,
-		})
-		if h.Len() > maxItem {
-			heap.Remove(h, h.Len()-1)
-		}
+	iter, err := c.db.NewIter(nil)
+	if err != nil {
+		c.log.Error("get iter failed", "error", err)
+		return
 	}
-	iterator.Release()
-	// delete the keys
-	for h.Len() > 0 {
-		if curentSize > expectSize {
+
+	batch := c.db.NewBatch()
+	for iter.Last(); iter.Valid(); iter.Prev() {
+		if curentSize < expectSize {
+			batch.Delete(iter.Key(), nil)
+			curentSize += uint64(len(iter.Key())) + uint64(len(iter.Value()))
+		} else {
+			distance := iter.Key()
+			c.db.Set(storage.RadisuKey, distance, c.writeOptions)
+			dis := uint256.NewInt(0)
+			err = dis.UnmarshalSSZ(distance)
+			if err != nil {
+				c.log.Error("unmarshal distance failed", "error", err)
+			}
+			c.radius.Store(dis)
 			break
 		}
-		item := heap.Pop(h)
-		val := item.(Item)
-		distance = val.Distance
-		key := xor(val.Distance, c.nodeId[:])
-		if err := c.db.Delete(key); err != nil {
-			c.log.Error("failed to delete key %v, err: %v", key, err)
-			continue
-		}
-		curentSize += val.ValueSize
 	}
-
-	dis := uint256.NewInt(0)
-	err := dis.UnmarshalSSZ(distance)
+	err = batch.Commit(&pebble.WriteOptions{Sync: true})
 	if err != nil {
-		c.log.Error("failed to parse the radius key %v, err is %v", distance, err)
-	}
-	c.radius.Store(dis)
-	err = c.db.Put(storage.RadisuKey, distance)
-
-	if err != nil {
-		c.log.Error("failed to save the radius key %v, err is %v", distance, err)
+		c.log.Error("prune batch commit failed", "error", err)
+		return
 	}
 }
 

portalnetwork/storage/ethpepple/storage_test.go

@@ -1,18 +1,14 @@
 package ethpepple
 
 import (
-	"crypto/rand"
 	"testing"
 	"time"
 
-	"github.com/ethereum/go-ethereum/ethdb/memorydb"
 	"github.com/ethereum/go-ethereum/portalnetwork/storage"
 	"github.com/holiman/uint256"
 	"github.com/stretchr/testify/assert"
 )
 
-var testRadius = uint256.NewInt(100000)
-
 func genBytes(length int) []byte {
 	res := make([]byte, length)
 	for i := 0; i < length; i++ {
@@ -21,68 +17,119 @@ func genBytes(length int) []byte {
 	return res
 }
 
-func getTestDb() (storage.ContentStorage, error) {
-	db := memorydb.New()
+func TestNewPeppleDB(t *testing.T) {
+	db, err := NewPeppleDB(t.TempDir(), 16, 16, "test")
+	assert.NoError(t, err)
+	defer db.Close()
+
+	assert.NotNil(t, db)
+}
+
+func setupTestStorage(t *testing.T) storage.ContentStorage {
+	db, err := NewPeppleDB(t.TempDir(), 16, 16, "test")
+	assert.NoError(t, err)
+	t.Cleanup(func() { db.Close() })
+
 	config := PeppleStorageConfig{
-		DB:                db,
 		StorageCapacityMB: 1,
+		DB:                db,
 		NodeId:            uint256.NewInt(0).Bytes32(),
-		NetworkName:       "history",
+		NetworkName:       "test",
 	}
-	return NewPeppleStorage(config)
+
+	storage, err := NewPeppleStorage(config)
+	assert.NoError(t, err)
+	return storage
 }
 
-func TestReadRadius(t *testing.T) {
-	db, err := getTestDb()
-	assert.NoError(t, err)
-	assert.True(t, db.Radius().Eq(storage.MaxDistance))
+func TestContentStoragePutAndGet(t *testing.T) {
+	db := setupTestStorage(t)
 
-	data, err := testRadius.MarshalSSZ()
-	assert.NoError(t, err)
-	db.Put(nil, storage.RadisuKey, data)
-
-	store := db.(*ContentStorage)
-	err = store.db.Close()
-	assert.NoError(t, err)
-}
-
-func TestStorage(t *testing.T) {
-	db, err := getTestDb()
-	assert.NoError(t, err)
-	testcases := map[string][]byte{
-		"test1": []byte("test1"),
-		"test2": []byte("test2"),
-		"test3": []byte("test3"),
-		"test4": []byte("test4"),
+	testCases := []struct {
+		contentKey []byte
+		contentId  []byte
+		content    []byte
+	}{
+		{[]byte("key1"), []byte("id1"), []byte("content1")},
+		{[]byte("key2"), []byte("id2"), []byte("content2")},
 	}
 
-	for key, value := range testcases {
-		db.Put(nil, []byte(key), value)
-	}
-
-	for key, value := range testcases {
-		val, err := db.Get(nil, []byte(key))
+	for _, tc := range testCases {
+		err := db.Put(tc.contentKey, tc.contentId, tc.content)
 		assert.NoError(t, err)
-		assert.Equal(t, value, val)
+
+		got, err := db.Get(tc.contentKey, tc.contentId)
+		assert.NoError(t, err)
+		assert.Equal(t, tc.content, got)
 	}
 }
 
-func TestXor(t *testing.T) {
-	nodeId := uint256.NewInt(0).Bytes32()
-	bs := make([]byte, 32)
-	rand.Read(bs)
-	dis := xor(bs, nodeId[:])
-	assert.Equal(t, bs, dis)
+func TestRadius(t *testing.T) {
+	db := setupTestStorage(t)
+	radius := db.Radius()
+	assert.NotNil(t, radius)
+	assert.True(t, radius.Eq(storage.MaxDistance))
+}
 
-	nodeId2 := uint256.NewInt(2).Bytes32()
-	dis = xor(bs, nodeId2[:])
-	assert.Equal(t, bs, xor(dis, nodeId2[:]))
+// func TestPrune(t *testing.T) {
+//     db, err := NewPeppleDB(t.TempDir(), 16, 16, "test")
+//     assert.NoError(t, err)
+//     defer db.Close()
+
+//     config := PeppleStorageConfig{
+//         StorageCapacityMB: 1, // 1MB capacity
+//         DB:                db,
+//         NodeId:            uint256.NewInt(0).Bytes32(),
+//         NetworkName:       "test",
+//     }
+
+//     storage, err := NewPeppleStorage(config)
+//     assert.NoError(t, err)
+
+//     // Add content exceeding capacity
+//     largeContent := make([]byte, 900_000) // 900KB
+//     err = storage.Put([]byte("key1"), []byte("id1"), largeContent)
+//     assert.NoError(t, err)
+
+//     smallContent := make([]byte, 200_000) // 200KB
+//     err = storage.Put([]byte("key2"), []byte("id2"), smallContent)
+//     assert.NoError(t, err)
+
+//     // Wait for prune to complete
+//     time.Sleep(6 * time.Second)
+
+//     // Verify content after pruning
+//     _, err = storage.Get([]byte("key2"), []byte("id2"))
+//     assert.Error(t, err) // Should be pruned
+// }
+
+func TestXOR(t *testing.T) {
+	testCases := []struct {
+		contentId []byte
+		nodeId    []byte
+		expected  []byte
+	}{
+		{
+			contentId: []byte{0x01},
+			nodeId:    make([]byte, 32),
+			expected:  append([]byte{0x01}, make([]byte, 31)...),
+		},
+		{
+			contentId: []byte{0xFF},
+			nodeId:    []byte{0x0F},
+			expected:  append([]byte{0xF0}, make([]byte, 31)...),
+		},
+	}
+
+	for _, tc := range testCases {
+		result := xor(tc.contentId, tc.nodeId)
+		assert.Equal(t, tc.expected, result)
+	}
 }
 
 // the capacity is 1MB, so prune will delete over 50Kb content
 func TestPrune(t *testing.T) {
-	db, err := getTestDb()
-	assert.NoError(t, err)
+	db := setupTestStorage(t)
 	// the nodeId is zeros, so contentKey and contentId is the same
 	testcases := []struct {
 		contentKey  [32]byte
@@ -130,7 +177,7 @@ func TestPrune(t *testing.T) {
 		db.Put(val.contentKey[:], val.contentKey[:], val.content)
 	}
 	// // wait to prune done
-	time.Sleep(5 * time.Second)
+	time.Sleep(2 * time.Second)
 	for _, val := range testcases {
 		content, err := db.Get(val.contentKey[:], val.contentKey[:])
 		if !val.shouldPrune {
@@ -139,4 +186,9 @@ func TestPrune(t *testing.T) {
 			assert.Error(t, err)
 		}
 	}
+	radius := db.Radius()
+	data, err := radius.MarshalSSZ()
+	assert.NoError(t, err)
+	actual := uint256.NewInt(4).Bytes32()
+	assert.Equal(t, data, actual)
 }