MT optimization: do not write nodes to DB until tx is commited (0xPol…

…ygonHermez#532)

* Changes for new MT spec and golden poseidon

* SC code hash working

* Add MT node cache

* Added rollback test

Makefile

@@ -153,6 +153,7 @@ generate-mocks: ## Generates mocks for the tests, using mockery tool
 	mockery --name=etherman --dir=sequencer/strategy/txprofitabilitychecker --output=sequencer/strategy/txprofitabilitychecker --outpkg=txprofitabilitychecker_test --filename=etherman-mock_test.go
 	mockery --name=batchProcessor --dir=sequencer/strategy/txselector --output=sequencer/strategy/txselector --outpkg=txselector_test --filename=batchprocessor-mock_test.go
 	mockery --name=etherman --dir=sequencer --output=sequencer --outpkg=sequencer --structname=ethermanMock --filename=etherman-mock_test.go
+	mockery --name=Store --dir=state/tree --output=state/tree --outpkg=tree --structname=storeMock --filename=store-mock_test.go
 
 .PHONY: generate-code-from-proto
 generate-code-from-proto: ## Generates code from proto files

state/tree/cache.go

@@ -1,6 +1,13 @@
 package tree
 
-import "github.com/dgraph-io/ristretto"
+import (
+	"errors"
+	"fmt"
+	"sync"
+
+	"github.com/dgraph-io/ristretto"
+	poseidon "github.com/iden3/go-iden3-crypto/goldenposeidon"
+)
 
 const (
 	cacheNumCounters = 1e7     // number of keys to track frequency of (10M).
@@ -21,3 +28,89 @@ func NewStoreCache() (*ristretto.Cache, error) {
 	}
 	return cache, err
 }
+
+// nodeCache represents a cache layer to store nodes during a db transaction in
+// the merkletree.
+type nodeCache struct {
+	lock   *sync.Mutex
+	data   map[string][]uint64
+	active bool
+}
+
+const (
+	maxMTNodeCacheEntries = 256
+)
+
+var (
+	errMTNodeCacheItemNotFound = errors.New("MT node cache item not found")
+)
+
+// newNodeCache is the nodeCache constructor.
+func newNodeCache() *nodeCache {
+	return &nodeCache{
+		lock: &sync.Mutex{},
+		data: make(map[string][]uint64),
+	}
+}
+
+// get reads a MT node cache entry.
+func (nc *nodeCache) get(key []uint64) ([]uint64, error) {
+	if len(key) != poseidon.CAPLEN {
+		return nil, fmt.Errorf("Invalid key length should be %d", poseidon.CAPLEN)
+	}
+	keyStr := h4ToString(key)
+
+	item, ok := nc.data[keyStr]
+	if !ok {
+		return nil, errMTNodeCacheItemNotFound
+	}
+	return item, nil
+}
+
+// set inserts a new MT node cache entry.
+func (nc *nodeCache) set(key []uint64, value []uint64) error {
+	if len(nc.data) >= maxMTNodeCacheEntries {
+		return errors.New("MT node cache is full")
+	}
+	if len(key) != poseidon.CAPLEN {
+		return fmt.Errorf("Invalid key length, should be %d", poseidon.CAPLEN)
+	}
+	keyStr := h4ToString(key)
+
+	nc.lock.Lock()
+	defer nc.lock.Unlock()
+
+	nc.data[keyStr] = value
+
+	return nil
+}
+
+// clear removes all the entries of the MT node cache.
+func (nc *nodeCache) clear() {
+	nc.lock.Lock()
+	defer nc.lock.Unlock()
+
+	nc.data = make(map[string][]uint64)
+}
+
+// isActive is the active field getter.
+func (nc *nodeCache) isActive() bool {
+	return nc.active
+}
+
+// setActive is the active field setter.
+func (nc *nodeCache) setActive(active bool) {
+	nc.active = active
+}
+
+// init initializes the MT node cache.
+func (nc *nodeCache) init() {
+	nc.clear()
+	nc.setActive(true)
+}
+
+// teardown resets the MT node cache.
+func (nc *nodeCache) teardown() {
+	nc.clear()
+	nc.setActive(false)
+}

state/tree/cache_test.go

@@ -0,0 +1,154 @@
+package tree
+
+import (
+	"testing"
+
+	"github.com/hermeznetwork/hermez-core/test/testutils"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMTNodeCacheGet(t *testing.T) {
+	tcs := []struct {
+		description    string
+		data           map[string][]uint64
+		key            [][]uint64
+		expected       [][]uint64
+		expectedErr    bool
+		expectedErrMsg string
+	}{
+		{
+			description: "single matching item",
+			data: map[string][]uint64{
+				"0x0000000000000001000000000000000100000000000000010000000000000001": {15},
+			},
+			key:      [][]uint64{{1, 1, 1, 1}},
+			expected: [][]uint64{{15}},
+		},
+		{
+			description: "single non-matching item",
+			data: map[string][]uint64{
+				"0x0000000000000001000000000000000100000000000000010000000000000001": {15},
+			},
+			key:            [][]uint64{{1, 1, 1, 0}},
+			expectedErr:    true,
+			expectedErrMsg: errMTNodeCacheItemNotFound.Error(),
+		},
+		{
+			description: "multiple matching items",
+			data: map[string][]uint64{
+				"0x0000000000000001000000000000000100000000000000010000000000000001": {15},
+				"0x0000000000000001000000000000000100000000000000010000000000000002": {16},
+				"0x0000000000000001000000000000000100000000000000010000000000000003": {17},
+			},
+			key:      [][]uint64{{1, 1, 1, 1}, {2, 1, 1, 1}, {3, 1, 1, 1}},
+			expected: [][]uint64{{15}, {16}, {17}},
+		},
+	}
+
+	for _, tc := range tcs {
+		tc := tc
+		t.Run(tc.description, func(t *testing.T) {
+			t.Parallel()
+
+			subject := newNodeCache()
+
+			subject.data = tc.data
+
+			for i := 0; i < len(tc.key); i++ {
+				actual, err := subject.get(tc.key[i])
+				require.NoError(t, testutils.CheckError(err, tc.expectedErr, tc.expectedErrMsg))
+
+				if !tc.expectedErr {
+					require.Equal(t, tc.expected[i], actual)
+				}
+			}
+		})
+	}
+}
+
+func TestMTNodeCacheSet(t *testing.T) {
+	tcs := []struct {
+		description    string
+		key            [][]uint64
+		value          [][]uint64
+		expectedData   map[string][]uint64
+		expectedErr    bool
+		expectedErrMsg string
+	}{
+		{
+			description: "single item set",
+			key:         [][]uint64{{1, 1, 1, 1}},
+			value:       [][]uint64{{15}},
+			expectedData: map[string][]uint64{
+				"0x0000000000000001000000000000000100000000000000010000000000000001": {15},
+			},
+		},
+		{
+			description: "mutiple items set",
+			key:         [][]uint64{{1, 1, 1, 1}, {1, 1, 1, 2}, {1, 2, 1, 3}},
+			value:       [][]uint64{{15}, {16}, {17}},
+			expectedData: map[string][]uint64{
+				"0x0000000000000001000000000000000100000000000000010000000000000001": {15},
+				"0x0000000000000002000000000000000100000000000000010000000000000001": {16},
+				"0x0000000000000003000000000000000100000000000000020000000000000001": {17},
+			},
+		},
+		{
+			description: "keys can be updated",
+			key:         [][]uint64{{1, 1, 1, 1}, {1, 1, 1, 2}, {1, 1, 1, 1}},
+			value:       [][]uint64{{15}, {16}, {1500}},
+			expectedData: map[string][]uint64{
+				"0x0000000000000001000000000000000100000000000000010000000000000001": {1500},
+				"0x0000000000000002000000000000000100000000000000010000000000000001": {16},
+			},
+		},
+		{
+			description:    "key length != 4 causes error",
+			key:            [][]uint64{{1, 1, 1}},
+			value:          [][]uint64{{15}},
+			expectedErr:    true,
+			expectedErrMsg: "Invalid key length, should be 4",
+		},
+	}
+
+	for _, tc := range tcs {
+		tc := tc
+		t.Run(tc.description, func(t *testing.T) {
+			t.Parallel()
+
+			subject := newNodeCache()
+			for i := 0; i < len(tc.key); i++ {
+				err := subject.set(tc.key[i], tc.value[i])
+				require.NoError(t, testutils.CheckError(err, tc.expectedErr, tc.expectedErrMsg))
+			}
+			if !tc.expectedErr {
+				require.Equal(t, tc.expectedData, subject.data)
+			}
+		})
+	}
+}
+
+func TestMTNodeCacheMaxItems(t *testing.T) {
+	subject := newNodeCache()
+	for i := 0; i < maxMTNodeCacheEntries; i++ {
+		err := subject.set([]uint64{1, 1, 1, uint64(i)}, []uint64{1})
+		require.NoError(t, err)
+	}
+
+	err := subject.set([]uint64{1, 1, 1, uint64(maxMTNodeCacheEntries + 1)}, []uint64{1})
+	require.Error(t, err)
+
+	require.Equal(t, "MT node cache is full", err.Error())
+}
+
+func TestMTNodeCacheClear(t *testing.T) {
+	subject := newNodeCache()
+	for i := 0; i < 5; i++ {
+		err := subject.set([]uint64{1, 1, 1, uint64(i)}, []uint64{1})
+		require.NoError(t, err)
+	}
+
+	subject.clear()
+
+	require.Zero(t, len(subject.data))
+}