storj/satellite/overlay/piececount_test.go

120 lines
2.8 KiB
Go
Raw Normal View History

// Copyright (C) 2019 Storj Labs, Inc.
// See LICENSE for copying information.
package overlay_test
import (
"math"
"testing"
"github.com/stretchr/testify/require"
"storj.io/common/pb"
"storj.io/common/storj"
"storj.io/common/testcontext"
"storj.io/common/testrand"
"storj.io/storj/satellite"
"storj.io/storj/satellite/overlay"
"storj.io/storj/satellite/satellitedb/satellitedbtest"
)
func TestDB_PieceCounts(t *testing.T) {
satellitedbtest.Run(t, func(t *testing.T, db satellite.DB) {
ctx := testcontext.New(t)
defer ctx.Cleanup()
overlaydb := db.OverlayCache()
type TestNode struct {
ID storj.NodeID
PieceCount int // TODO: fix to int64
}
nodes := make([]TestNode, 100)
for i := range nodes {
nodes[i].ID = testrand.NodeID()
nodes[i].PieceCount = int(math.Pow10(i + 1))
}
for _, node := range nodes {
require.NoError(t, overlaydb.UpdateAddress(ctx, &pb.Node{
Id: node.ID,
Address: &pb.NodeAddress{
Transport: pb.NodeTransport_TCP_TLS_GRPC,
Address: "0.0.0.0",
},
LastIp: "0.0.0.0",
}, overlay.NodeSelectionConfig{}))
}
// check that they are initialized to zero
initialCounts, err := overlaydb.AllPieceCounts(ctx)
require.NoError(t, err)
require.Empty(t, initialCounts)
// TODO: make AllPieceCounts return results for all nodes,
// since it will keep the logic slightly clearer.
// update counts
counts := make(map[storj.NodeID]int)
for _, node := range nodes {
counts[node.ID] = node.PieceCount
}
err = overlaydb.UpdatePieceCounts(ctx, counts)
require.NoError(t, err)
// fetch new counts
updatedCounts, err := overlaydb.AllPieceCounts(ctx)
require.NoError(t, err)
// verify values
for _, node := range nodes {
count, ok := updatedCounts[node.ID]
require.True(t, ok)
require.Equal(t, count, node.PieceCount)
}
})
}
func BenchmarkDB_PieceCounts(b *testing.B) {
satellitedbtest.Bench(b, func(b *testing.B, db satellite.DB) {
ctx := testcontext.New(b)
defer ctx.Cleanup()
overlaydb := db.OverlayCache()
counts := make(map[storj.NodeID]int)
for i := 0; i < 10000; i++ {
counts[testrand.NodeID()] = testrand.Intn(100000)
}
for nodeID := range counts {
require.NoError(b, overlaydb.UpdateAddress(ctx, &pb.Node{
Id: nodeID,
Address: &pb.NodeAddress{
Transport: pb.NodeTransport_TCP_TLS_GRPC,
Address: "0.0.0.0",
},
LastIp: "0.0.0.0",
}, overlay.NodeSelectionConfig{}))
}
b.Run("Update", func(b *testing.B) {
for i := 0; i < b.N; i++ {
err := overlaydb.UpdatePieceCounts(ctx, counts)
if err != nil {
b.Fatal(err)
}
}
})
b.Run("All", func(b *testing.B) {
for i := 0; i < b.N; i++ {
_, err := overlaydb.AllPieceCounts(ctx)
if err != nil {
b.Fatal(err)
}
}
})
})
}