feat: delete content and set radius

portalnetwork/storage/ethpepple/maxheap.go

@@ -0,0 +1,39 @@
+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

@@ -0,0 +1,33 @@
+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,7 +1,12 @@
 package ethpepple
 
 import (
+	"bytes"
+	"container/heap"
+	"encoding/binary"
+	"sync"
 	"sync/atomic"
+	"time"
 
 	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/ethdb/pebble"
@@ -11,6 +16,8 @@ import (
 	"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.
+
 var _ storage.ContentStorage = &ContentStorage{}
 
 type PeppleStorageConfig struct {
@@ -29,9 +36,13 @@ type ContentStorage struct {
 	nodeId                 enode.ID
 	storageCapacityInBytes uint64
 	radius                 atomic.Value
-	// size 									 uint64
-	log log.Logger
-	db  ethdb.KeyValueStore
+	log                    log.Logger
+	db                     ethdb.KeyValueStore
+	size                   uint64
+	sizeChan               chan uint64
+	sizeMutex              sync.RWMutex
+	isPruning              bool
+	pruneDoneChan          chan uint64 // finish prune and get the pruned size
 }
 
 func NewPeppleStorage(config PeppleStorageConfig) (storage.ContentStorage, error) {
@@ -40,6 +51,8 @@ func NewPeppleStorage(config PeppleStorageConfig) (storage.ContentStorage, error
 		db:                     config.DB,
 		storageCapacityInBytes: config.StorageCapacityMB * 1000_000,
 		log:                    log.New("storage", config.NetworkName),
+		sizeChan:               make(chan uint64, 100),
+		pruneDoneChan:          make(chan uint64, 1),
 	}
 	cs.radius.Store(storage.MaxDistance)
 	exist, err := cs.db.Has(storage.RadisuKey)
@@ -58,6 +71,21 @@ func NewPeppleStorage(config PeppleStorageConfig) (storage.ContentStorage, error
 		}
 		cs.radius.Store(dis)
 	}
+
+	exist, err = cs.db.Has(storage.SizeKey)
+	if err != nil {
+		return nil, err
+	}
+	if exist {
+		val, err := cs.db.Get(storage.SizeKey)
+		if err != nil {
+			return nil, err
+		}
+		size := binary.BigEndian.Uint64(val)
+		// init stage, no need to use lock
+		cs.size = size
+	}
+	go cs.saveCapacity()
 	return cs, nil
 }
 
@@ -68,6 +96,8 @@ func (c *ContentStorage) Get(contentKey []byte, contentId []byte) ([]byte, error
 
 // 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)
 }
 
@@ -77,3 +107,115 @@ func (c *ContentStorage) Radius() *uint256.Int {
 	val := radius.(*uint256.Int)
 	return val
 }
+
+func (c *ContentStorage) saveCapacity() {
+	ticker := time.NewTicker(5 * time.Second)
+	defer ticker.Stop()
+	sizeChanged := false
+	buf := make([]byte, 8) // uint64
+
+	for {
+		select {
+		case <-ticker.C:
+			if sizeChanged {
+				binary.BigEndian.PutUint64(buf, c.size)
+				c.db.Put(storage.SizeKey, buf)
+				sizeChanged = false
+			}
+		case size := <-c.sizeChan:
+			c.log.Debug("reveice size %v", size)
+			c.sizeMutex.Lock()
+			c.size += size
+			c.sizeMutex.Unlock()
+			sizeChanged = true
+			if c.size > c.storageCapacityInBytes {
+				if !c.isPruning {
+					c.isPruning = true
+					go c.prune()
+				}
+			}
+		case prunedSize := <-c.pruneDoneChan:
+			c.isPruning = false
+			c.size -= prunedSize
+			sizeChanged = true
+		}
+	}
+}
+
+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)
+		}
+	}
+	iterator.Release()
+	// delete the keys
+	for h.Len() > 0 {
+		if curentSize > expectSize {
+			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)
+	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)
+	}
+}
+
+func xor(contentId, nodeId []byte) []byte {
+	// length of contentId maybe not 32bytes
+	padding := make([]byte, 32)
+	if len(contentId) != len(nodeId) {
+		copy(padding, contentId)
+	} else {
+		padding = contentId
+	}
+	res := make([]byte, len(padding))
+	for i := range padding {
+		res[i] = padding[i] ^ nodeId[i]
+	}
+	return res
+}

portalnetwork/storage/ethpepple/storage_test.go

@@ -2,11 +2,10 @@ package ethpepple
 
 import (
 	"crypto/rand"
-	"encoding/hex"
-	"os"
 	"testing"
+	"time"
 
-	"github.com/ethereum/go-ethereum/p2p/enode"
+	"github.com/ethereum/go-ethereum/ethdb/memorydb"
 	"github.com/ethereum/go-ethereum/portalnetwork/storage"
 	"github.com/holiman/uint256"
 	"github.com/stretchr/testify/assert"
@@ -14,39 +13,28 @@ import (
 
 var testRadius = uint256.NewInt(100000)
 
-func clearNodeData(path string) {
-	_ = os.RemoveAll(path)
+func genBytes(length int) []byte {
+	res := make([]byte, length)
+	for i := 0; i < length; i++ {
+		res[i] = byte(i)
+	}
+	return res
 }
 
-func getRandomPath() string {
-	// gen a random hex string
-	bytes := make([]byte, 32)
-	_, err := rand.Read(bytes)
-	if err != nil {
-		panic(err)
-	}
-	return hex.EncodeToString(bytes)
-}
-
-func getTestDb(path string) (storage.ContentStorage, error) {
-	db, err := NewPeppleDB(path, 100, 100, "history")
-	if err != nil {
-		return nil, err
-	}
+func getTestDb() (storage.ContentStorage, error) {
+	db := memorydb.New()
 	config := PeppleStorageConfig{
 		DB:                db,
-		StorageCapacityMB: 100,
-		NodeId:            enode.ID{},
+		StorageCapacityMB: 1,
+		NodeId:            uint256.NewInt(0).Bytes32(),
 		NetworkName:       "history",
 	}
 	return NewPeppleStorage(config)
 }
 
 func TestReadRadius(t *testing.T) {
-	path := getRandomPath()
-	db, err := getTestDb(path)
+	db, err := getTestDb()
 	assert.NoError(t, err)
-	defer clearNodeData(path)
 	assert.True(t, db.Radius().Eq(storage.MaxDistance))
 
 	data, err := testRadius.MarshalSSZ()
@@ -56,17 +44,11 @@ func TestReadRadius(t *testing.T) {
 	store := db.(*ContentStorage)
 	err = store.db.Close()
 	assert.NoError(t, err)
-
-	db, err = getTestDb(path)
-	assert.NoError(t, err)
-	assert.True(t, db.Radius().Eq(testRadius))
 }
 
 func TestStorage(t *testing.T) {
-	path := getRandomPath()
-	db, err := getTestDb(path)
+	db, err := getTestDb()
 	assert.NoError(t, err)
-	defer clearNodeData(path)
 	testcases := map[string][]byte{
 		"test1": []byte("test1"),
 		"test2": []byte("test2"),
@@ -84,3 +66,77 @@ func TestStorage(t *testing.T) {
 		assert.Equal(t, value, val)
 	}
 }
+
+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)
+
+	nodeId2 := uint256.NewInt(2).Bytes32()
+	dis = xor(bs, nodeId2[:])
+	assert.Equal(t, bs, xor(dis, nodeId2[:]))
+}
+
+// the capacity is 1MB, so prune will delete over 50Kb content
+func TestPrune(t *testing.T) {
+	db, err := getTestDb()
+	assert.NoError(t, err)
+	// the nodeId is zeros, so contentKey and contentId is the same
+	testcases := []struct {
+		contentKey  [32]byte
+		content     []byte
+		shouldPrune bool
+	}{
+		{
+			contentKey:  uint256.NewInt(1).Bytes32(),
+			content:     genBytes(900_000),
+			shouldPrune: false,
+		},
+		{
+			contentKey:  uint256.NewInt(2).Bytes32(),
+			content:     genBytes(40_000),
+			shouldPrune: false,
+		},
+		{
+			contentKey:  uint256.NewInt(3).Bytes32(),
+			content:     genBytes(20_000),
+			shouldPrune: false,
+		},
+		{
+			contentKey:  uint256.NewInt(4).Bytes32(),
+			content:     genBytes(20_000),
+			shouldPrune: false,
+		},
+		{
+			contentKey:  uint256.NewInt(5).Bytes32(),
+			content:     genBytes(20_000),
+			shouldPrune: true,
+		},
+		{
+			contentKey:  uint256.NewInt(6).Bytes32(),
+			content:     genBytes(20_000),
+			shouldPrune: true,
+		},
+		{
+			contentKey:  uint256.NewInt(7).Bytes32(),
+			content:     genBytes(20_000),
+			shouldPrune: true,
+		},
+	}
+
+	for _, val := range testcases {
+		db.Put(val.contentKey[:], val.contentKey[:], val.content)
+	}
+	// // wait to prune done
+	time.Sleep(5 * time.Second)
+	for _, val := range testcases {
+		content, err := db.Get(val.contentKey[:], val.contentKey[:])
+		if !val.shouldPrune {
+			assert.Equal(t, val.content, content)
+		} else {
+			assert.Error(t, err)
+		}
+	}
+}