Merge pull request #2057 from nspcc-dev/optimize-storage-find

Optimize `storageFind`

pkg/core/helper_test.go

@@ -41,6 +41,7 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/vm/stackitem"
 	"github.com/nspcc-dev/neo-go/pkg/wallet"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
 	"go.uber.org/zap/zaptest"
 )
 
@@ -74,22 +75,22 @@ func newTestChainWithNS(t *testing.T) (*Blockchain, util.Uint160) {
 
 // newTestChain should be called before newBlock invocation to properly setup
 // global state.
-func newTestChain(t *testing.T) *Blockchain {
+func newTestChain(t testing.TB) *Blockchain {
 	return newTestChainWithCustomCfg(t, nil)
 }
 
-func newTestChainWithCustomCfg(t *testing.T, f func(*config.Config)) *Blockchain {
+func newTestChainWithCustomCfg(t testing.TB, f func(*config.Config)) *Blockchain {
 	return newTestChainWithCustomCfgAndStore(t, nil, f)
 }
 
-func newTestChainWithCustomCfgAndStore(t *testing.T, st storage.Store, f func(*config.Config)) *Blockchain {
+func newTestChainWithCustomCfgAndStore(t testing.TB, st storage.Store, f func(*config.Config)) *Blockchain {
 	chain := initTestChain(t, st, f)
 	go chain.Run()
 	t.Cleanup(chain.Close)
 	return chain
 }
 
-func initTestChain(t *testing.T, st storage.Store, f func(*config.Config)) *Blockchain {
+func initTestChain(t testing.TB, st storage.Store, f func(*config.Config)) *Blockchain {
 	unitTestNetCfg, err := config.Load("../../config", testchain.Network())
 	require.NoError(t, err)
 	if f != nil {
@@ -98,7 +99,11 @@ func initTestChain(t *testing.T, st storage.Store, f func(*config.Config)) *Bloc
 	if st == nil {
 		st = storage.NewMemoryStore()
 	}
-	chain, err := NewBlockchain(st, unitTestNetCfg.ProtocolConfiguration, zaptest.NewLogger(t))
+	log := zaptest.NewLogger(t)
+	if _, ok := t.(*testing.B); ok {
+		log = zap.NewNop()
+	}
+	chain, err := NewBlockchain(st, unitTestNetCfg.ProtocolConfiguration, log)
 	require.NoError(t, err)
 	return chain
 }

pkg/core/interop_system.go

@@ -192,18 +192,23 @@ func storageFind(ic *interop.Context) error {
 		return err
 	}
 
-	filteredMap := stackitem.NewMap()
+	arr := make([]stackitem.MapElement, 0, len(siMap))
 	for k, v := range siMap {
-		key := append(prefix, []byte(k)...)
-		keycopy := make([]byte, len(key))
-		copy(keycopy, key)
-		filteredMap.Add(stackitem.NewByteArray(keycopy), stackitem.NewByteArray(v))
+		keycopy := make([]byte, len(k)+len(prefix))
+		copy(keycopy, prefix)
+		copy(keycopy[len(prefix):], k)
+		arr = append(arr, stackitem.MapElement{
+			Key:   stackitem.NewByteArray(keycopy),
+			Value: stackitem.NewByteArray(v),
+		})
 	}
-	sort.Slice(filteredMap.Value().([]stackitem.MapElement), func(i, j int) bool {
-		return bytes.Compare(filteredMap.Value().([]stackitem.MapElement)[i].Key.Value().([]byte),
-			filteredMap.Value().([]stackitem.MapElement)[j].Key.Value().([]byte)) == -1
+	sort.Slice(arr, func(i, j int) bool {
+		k1 := arr[i].Key.Value().([]byte)
+		k2 := arr[j].Key.Value().([]byte)
+		return bytes.Compare(k1, k2) == -1
 	})
 
+	filteredMap := stackitem.NewMapWithValue(arr)
 	item := istorage.NewIterator(filteredMap, len(prefix), opts)
 	ic.VM.Estack().PushVal(stackitem.NewInterop(item))
 

pkg/core/interop_system_test.go

@@ -259,6 +259,36 @@ func TestStorageDelete(t *testing.T) {
 	})
 }
 
+func BenchmarkStorageFind(b *testing.B) {
+	v, contractState, context, chain := createVMAndContractState(b)
+	require.NoError(b, chain.contracts.Management.PutContractState(chain.dao, contractState))
+
+	const count = 100
+
+	items := make(map[string]state.StorageItem)
+	for i := 0; i < count; i++ {
+		items["abc"+random.String(10)] = random.Bytes(10)
+	}
+	for k, v := range items {
+		require.NoError(b, context.DAO.PutStorageItem(contractState.ID, []byte(k), v))
+		require.NoError(b, context.DAO.PutStorageItem(contractState.ID+1, []byte(k), v))
+	}
+
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		v.Estack().PushVal(istorage.FindDefault)
+		v.Estack().PushVal("abc")
+		v.Estack().PushVal(stackitem.NewInterop(&StorageContext{ID: contractState.ID}))
+		b.StartTimer()
+		err := storageFind(context)
+		if err != nil {
+			b.FailNow()
+		}
+	}
+}
+
 func TestStorageFind(t *testing.T) {
 	v, contractState, context, chain := createVMAndContractState(t)
 
@@ -451,7 +481,7 @@ func createVM(t *testing.T) (*vm.VM, *interop.Context, *Blockchain) {
 	return v, context, chain
 }
 
-func createVMAndContractState(t *testing.T) (*vm.VM, *state.Contract, *interop.Context, *Blockchain) {
+func createVMAndContractState(t testing.TB) (*vm.VM, *state.Contract, *interop.Context, *Blockchain) {
 	script := []byte("testscript")
 	m := manifest.NewManifest("Test")
 	ne, err := nef.NewFile(script)

pkg/core/storage/memory_store.go

@@ -121,8 +121,9 @@ func (s *MemoryStore) SeekAll(key []byte, f func(k, v []byte)) {
 
 // seek is an internal unlocked implementation of Seek.
 func (s *MemoryStore) seek(key []byte, f func(k, v []byte)) {
+	sk := string(key)
 	for k, v := range s.mem {
-		if strings.HasPrefix(k, string(key)) {
+		if strings.HasPrefix(k, sk) {
 			f([]byte(k), v)
 		}
 	}