Merge pull request #1638 from nspcc-dev/fix/mpt

mpt: do not allocate new buffer when updating dirty node

pkg/core/blockchain.go

@@ -687,7 +687,12 @@ func (bc *Blockchain) storeBlock(block *block.Block, txpool *mempool.Pool) error
 	}
 	writeBuf.Reset()
 
-	root := bc.dao.MPT.StateRoot()
+	d := cache.DAO.(*dao.Simple)
+	if err := d.UpdateMPT(); err != nil {
+		return fmt.Errorf("error while trying to apply MPT changes: %w", err)
+	}
+
+	root := d.MPT.StateRoot()
 	var prevHash util.Uint256
 	if block.Index > 0 {
 		prev, err := bc.dao.GetStateRoot(block.Index - 1)

pkg/core/dao/dao.go

@@ -382,11 +382,6 @@ func (dao *Simple) PutStorageItem(id int32, key []byte, si *state.StorageItem) e
 		return buf.Err
 	}
 	v := buf.Bytes()
-	if dao.MPT != nil {
-		if err := dao.MPT.Put(stKey[1:], v); err != nil && err != mpt.ErrNotFound {
-			return err
-		}
-	}
 	return dao.Store.Put(stKey, v)
 }
 
@@ -394,11 +389,6 @@ func (dao *Simple) PutStorageItem(id int32, key []byte, si *state.StorageItem) e
 // given key from the store.
 func (dao *Simple) DeleteStorageItem(id int32, key []byte) error {
 	stKey := makeStorageItemKey(id, key)
-	if dao.MPT != nil {
-		if err := dao.MPT.Delete(stKey[1:]); err != nil && err != mpt.ErrNotFound {
-			return err
-		}
-	}
 	return dao.Store.Delete(stKey)
 }
 
@@ -701,3 +691,18 @@ func (dao *Simple) StoreAsTransaction(tx *transaction.Transaction, index uint32,
 func (dao *Simple) Persist() (int, error) {
 	return dao.Store.Persist()
 }
+
+// UpdateMPT updates MPT using storage items from the underlying memcached store.
+func (dao *Simple) UpdateMPT() error {
+	var err error
+	dao.Store.MemoryStore.SeekAll([]byte{byte(storage.STStorage)}, func(k, v []byte) {
+		if err != nil {
+			return
+		} else if v != nil {
+			err = dao.MPT.Put(k[1:], v)
+		} else {
+			err = dao.MPT.Delete(k[1:])
+		}
+	})
+	return err
+}

pkg/core/mpt/base.go

@@ -63,7 +63,7 @@ func (b *BaseNode) updateHash(n Node) {
 
 // updateCache updates hash and bytes fields for this BaseNode.
 func (b *BaseNode) updateBytes(n Node) {
-	buf := io.NewBufBinWriter()
+	buf := io.NewBufBinWriterPreAlloc(b.bytes)
 	encodeNodeWithType(n, buf.BinWriter)
 	b.bytes = buf.Bytes()
 	b.bytesValid = true

pkg/core/mpt/trie.go

@@ -418,31 +418,29 @@ func (t *Trie) updateRefCount(h util.Uint256) int32 {
 func (t *Trie) addRef(h util.Uint256, bs []byte) {
 	node := t.refcount[h]
 	if node == nil {
+		data := make([]byte, len(bs))
+		copy(data, bs)
 		t.refcount[h] = &cachedNode{
 			refcount: 1,
-			bytes:    bs,
+			bytes:    data,
 		}
 		return
 	}
 	node.refcount++
-	if node.bytes == nil {
-		node.bytes = bs
-	}
 }
 
 func (t *Trie) removeRef(h util.Uint256, bs []byte) {
 	node := t.refcount[h]
 	if node == nil {
+		data := make([]byte, len(bs))
+		copy(data, bs)
 		t.refcount[h] = &cachedNode{
 			refcount: -1,
-			bytes:    bs,
+			bytes:    data,
 		}
 		return
 	}
 	node.refcount--
-	if node.bytes == nil {
-		node.bytes = bs
-	}
 }
 
 func (t *Trie) getFromStore(h util.Uint256) (Node, error) {
@@ -462,7 +460,8 @@ func (t *Trie) getFromStore(h util.Uint256) (Node, error) {
 		data = data[:len(data)-4]
 		node := t.refcount[h]
 		if node != nil {
-			node.bytes = data
+			node.bytes = make([]byte, len(data))
+			copy(node.bytes, data)
 			node.initial = int32(r.ReadU32LE())
 		}
 	}

pkg/core/native_management_test.go

@@ -440,6 +440,7 @@ func TestContractDestroy(t *testing.T) {
 	require.NoError(t, err)
 	err = bc.dao.PutStorageItem(cs1.ID, []byte{1, 2, 3}, &state.StorageItem{Value: []byte{3, 2, 1}})
 	require.NoError(t, err)
+	require.NoError(t, bc.dao.UpdateMPT())
 
 	t.Run("no contract", func(t *testing.T) {
 		res, err := invokeContractMethod(bc, 1_00000000, mgmtHash, "destroy")

pkg/core/storage/memory_store.go

@@ -102,6 +102,23 @@ func (s *MemoryStore) Seek(key []byte, f func(k, v []byte)) {
 	s.mut.RUnlock()
 }
 
+// SeekAll is like seek but also iterates over deleted items.
+func (s *MemoryStore) SeekAll(key []byte, f func(k, v []byte)) {
+	s.mut.RLock()
+	defer s.mut.RUnlock()
+	sk := string(key)
+	for k, v := range s.mem {
+		if strings.HasPrefix(k, sk) {
+			f([]byte(k), v)
+		}
+	}
+	for k := range s.del {
+		if strings.HasPrefix(k, sk) {
+			f([]byte(k), nil)
+		}
+	}
+}
+
 // seek is an internal unlocked implementation of Seek.
 func (s *MemoryStore) seek(key []byte, f func(k, v []byte)) {
 	for k, v := range s.mem {

pkg/io/binaryBufWriter.go

@@ -19,6 +19,12 @@ func NewBufBinWriter() *BufBinWriter {
 	return &BufBinWriter{BinWriter: NewBinWriterFromIO(b), buf: b}
 }
 
+// NewBufBinWriterPreAlloc makes a BufBinWriter using preallocated buffer.
+func NewBufBinWriterPreAlloc(buf []byte) *BufBinWriter {
+	b := bytes.NewBuffer(buf[:0])
+	return &BufBinWriter{BinWriter: NewBinWriterFromIO(b), buf: b}
+}
+
 // Len returns the number of bytes of the unread portion of the buffer.
 func (bw *BufBinWriter) Len() int {
 	return bw.buf.Len()