Implement hyper cache rebuild

Co-authored-by: Gabriel Díaz <[email protected]>

balloon/hyper2/navigation/position.go

@@ -26,13 +26,13 @@ type Position struct {
 	Index  []byte
 	Height uint16
 
-	serialized []byte
+	serialized []byte // height:index
 	numBits    uint16
 }
 
 func NewPosition(index []byte, height uint16) Position {
 	// Size of the index plus 2 bytes for the height
-	b := append(append([]byte{}, index[:len(index)]...), util.Uint16AsBytes(height)...)
+	b := append(util.Uint16AsBytes(height), index[:len(index)]...)
 	return Position{
 		Index:      index,
 		Height:     height,
@@ -41,8 +41,8 @@ func NewPosition(index []byte, height uint16) Position {
 	}
 }
 
-func NewRootPosition(numBits uint16) Position {
-	return NewPosition(make([]byte, numBits/8), numBits)
+func NewRootPosition(indexNumBytes uint16) Position {
+	return NewPosition(make([]byte, indexNumBytes), indexNumBytes*8)
 }
 
 func (p Position) Bytes() []byte {

balloon/hyper2/navigation/position_test.go

@@ -25,19 +25,19 @@ import (
 func TestRoot(t *testing.T) {
 
 	testCases := []struct {
-		numBits     uint16
-		expectedPos Position
+		indexNumBytes uint16
+		expectedPos   Position
 	}{
-		{8, NewPosition(make([]byte, 1), 8)},
-		{16, NewPosition(make([]byte, 2), 16)},
-		{32, NewPosition(make([]byte, 4), 32)},
-		{64, NewPosition(make([]byte, 8), 64)},
-		{128, NewPosition(make([]byte, 16), 128)},
-		{256, NewPosition(make([]byte, 32), 256)},
+		{1, NewPosition(make([]byte, 1), 8)},
+		{2, NewPosition(make([]byte, 2), 16)},
+		{4, NewPosition(make([]byte, 4), 32)},
+		{8, NewPosition(make([]byte, 8), 64)},
+		{16, NewPosition(make([]byte, 16), 128)},
+		{32, NewPosition(make([]byte, 32), 256)},
 	}
 
 	for i, c := range testCases {
-		rootPos := NewRootPosition(c.numBits)
+		rootPos := NewRootPosition(c.indexNumBytes)
 		require.Equalf(t, c.expectedPos, rootPos, "The root position should match in test case %d", i)
 	}
 

balloon/hyper2/navigation/test_util.go

@@ -16,8 +16,6 @@
 
 package navigation
 
-import "github.com/bbva/qed/util"
-
-func pos(index uint64, height uint16) Position {
-	return NewPosition(util.Uint64AsBytes(index), height)
+func pos(index uint8, height uint16) Position {
+	return NewPosition([]byte{byte(index)}, height)
 }

balloon/hyper2/pruning/insert.go

@@ -207,6 +207,6 @@ func PruneToInsert(index []byte, value []byte, cacheHeightLimit uint16, batches
 	ops := NewOperationsStack()
 	leaves := make(Leaves, 0)
 	leaves = leaves.InsertSorted(Leaf{index, value})
-	traverse(navigation.NewRootPosition(uint16(len(index)*8)), leaves, nil, 0, ops)
+	traverse(navigation.NewRootPosition(uint16(len(index))), leaves, nil, 0, ops)
 	return ops
 }

balloon/hyper2/pruning/operation.go

@@ -45,6 +45,7 @@ const (
 	MutateBatchCode
 	CollectHashCode
 	GetFromPathCode
+	UseHashCode
 	NoOpCode
 )
 
@@ -172,6 +173,16 @@ func getFromPath(pos navigation.Position) *Operation {
 	}
 }
 
+func useHash(pos navigation.Position, digest []byte) *Operation {
+	return &Operation{
+		Code: UseHashCode,
+		Pos:  pos,
+		Interpret: func(ops *OperationsStack, c *Context) hashing.Digest {
+			return digest
+		},
+	}
+}
+
 func noOp(pos navigation.Position) *Operation {
 	return &Operation{
 		Code: NoOpCode,

balloon/hyper2/pruning/rebuild.go

@@ -0,0 +1,81 @@
+/*
+   Copyright 2018 Banco Bilbao Vizcaya Argentaria, S.A.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package pruning
+
+import (
+	"bytes"
+
+	"github.com/bbva/qed/balloon/hyper2/navigation"
+)
+
+func PruneToRebuild(index, serializedBatch []byte, cacheHeightLimit uint16, batches BatchLoader) *OperationsStack {
+
+	persistedBatch := ParseBatchNode(len(index), serializedBatch)
+
+	var traverse, discardBranch func(pos navigation.Position, batch *BatchNode, iBatch int8, ops *OperationsStack)
+
+	discardBranch = func(pos navigation.Position, batch *BatchNode, iBatch int8, ops *OperationsStack) {
+
+		if batch.HasElementAt(iBatch) {
+			ops.Push(getProvidedHash(pos, iBatch, batch))
+		} else {
+			ops.Push(getDefaultHash(pos))
+		}
+	}
+
+	traverse = func(pos navigation.Position, batch *BatchNode, iBatch int8, ops *OperationsStack) {
+
+		// we don't need to check the length of the leaves because we
+		// always have to descend to the cache height limit
+		if pos.Height == cacheHeightLimit {
+			ops.PushAll(useHash(pos, persistedBatch.GetElementAt(0)), updateBatchNode(pos, iBatch, batch))
+			return
+		}
+
+		if batch == nil {
+			batch = batches.Load(pos)
+		}
+
+		// at the end of a batch tree
+		if iBatch > 0 && pos.Height%4 == 0 {
+			traverse(pos, nil, 0, ops)
+			ops.Push(updateBatchNode(pos, iBatch, batch))
+			return
+		}
+
+		rightPos := pos.Right()
+		leftPos := pos.Left()
+		if bytes.Compare(index, rightPos.Index) < 0 { // go to left
+			traverse(pos.Left(), batch, 2*iBatch+1, ops)
+			discardBranch(rightPos, batch, 2*iBatch+2, ops)
+		} else { // go to right
+			discardBranch(leftPos, batch, 2*iBatch+1, ops)
+			traverse(rightPos, batch, 2*iBatch+2, ops)
+		}
+
+		ops.PushAll(innerHash(pos), updateBatchNode(pos, iBatch, batch))
+		if iBatch == 0 { // it's the root of the batch tree
+			ops.Push(putInCache(pos, batch))
+		}
+
+	}
+
+	ops := NewOperationsStack()
+	traverse(navigation.NewRootPosition(uint16(len(index))), nil, 0, ops)
+	return ops
+
+}

balloon/hyper2/pruning/search.go

@@ -80,7 +80,7 @@ func PruneToFind(index []byte, batches BatchLoader) *OperationsStack {
 	}
 
 	ops := NewOperationsStack()
-	root := navigation.NewRootPosition(uint16(len(index) * 8))
+	root := navigation.NewRootPosition(uint16(len(index)))
 	traverse(root, nil, 0, ops)
 	if ops.Len() == 0 {
 		ops.Push(noOp(root))

balloon/hyper2/pruning/verify.go

@@ -47,7 +47,7 @@ func PruneToVerify(index, value []byte, auditPathHeight uint16) *OperationsStack
 	}
 
 	ops := NewOperationsStack()
-	traverse(navigation.NewRootPosition(uint16(len(index)*8)), ops)
+	traverse(navigation.NewRootPosition(uint16(len(index))), ops)
 	return ops
 
 }

balloon/hyper2/tree.go

@@ -20,6 +20,7 @@ import (
 	"sync"
 
 	"github.com/bbva/qed/balloon/hyper2/navigation"
+	"github.com/bbva/qed/log"
 
 	"github.com/bbva/qed/balloon/cache"
 	"github.com/bbva/qed/balloon/hyper2/pruning"
@@ -63,7 +64,7 @@ func NewHyperTree(hasherF func() hashing.Hasher, store storage.Store, cache cach
 	}
 
 	// warm-up cache
-	//tree.RebuildCache()
+	tree.RebuildCache()
 
 	return tree
 }
@@ -117,6 +118,43 @@ func (t *HyperTree) QueryMembership(eventDigest hashing.Digest, version []byte)
 	return NewQueryProof(eventDigest, version, ctx.AuditPath, t.hasherF()), nil
 }
 
+func (t *HyperTree) RebuildCache() {
+	t.Lock()
+	defer t.Unlock()
+
+	// warm up cache
+	log.Info("Warming up hyper cache...")
+
+	// get all nodes at cache limit height
+	start := make([]byte, 2+t.hasher.Len()/8)
+	end := make([]byte, 2+t.hasher.Len()/8)
+	start[1] = byte(t.cacheHeightLimit)
+	end[1] = byte(t.cacheHeightLimit + 1)
+	nodes, err := t.store.GetRange(storage.HyperCachePrefix, start, end)
+	if err != nil {
+		log.Fatalf("Oops, something went wrong: %v", err)
+	}
+
+	// insert every node into cache
+	for _, node := range nodes {
+		ops := pruning.PruneToRebuild(node.Key[2:], node.Value, t.cacheHeightLimit, t.batchLoader)
+		ctx := &pruning.Context{
+			Hasher:        t.hasher,
+			Cache:         t.cache,
+			DefaultHashes: t.defaultHashes,
+		}
+		ops.Pop().Interpret(ops, ctx)
+	}
+}
+
+func (t *HyperTree) Close() {
+	t.cache = nil
+	t.hasher = nil
+	t.defaultHashes = nil
+	t.store = nil
+	t.batchLoader = nil
+}
+
 func min(x, y uint16) uint16 {
 	if x < y {
 		return x

balloon/hyper2/tree_test.go

@@ -254,6 +254,36 @@ func TestDeterministicAdd(t *testing.T) {
 
 }
 
+func TestRebuildCache(t *testing.T) {
+
+	log.SetLogger("TestRebuildCache", log.SILENT)
+
+	store, closeF := storage_utils.OpenBPlusTreeStore()
+	defer closeF()
+	hasherF := hashing.NewSha256Hasher
+	hasher := hasherF()
+
+	firstCache := cache.NewSimpleCache(10)
+	tree := NewHyperTree(hasherF, store, firstCache)
+	require.True(t, firstCache.Size() == 0, "The cache should be empty")
+
+	// store multiple elements
+	for i := 0; i < 1000; i++ {
+		key := hasher.Do(rand.Bytes(32))
+		_, mutations, _ := tree.Add(key, uint64(i))
+		store.Mutate(mutations)
+	}
+	expectedSize := firstCache.Size()
+
+	// Close tree and reopen with a new fresh cache
+	tree.Close()
+	secondCache := cache.NewSimpleCache(10)
+	tree = NewHyperTree(hasherF, store, secondCache)
+
+	require.Equal(t, expectedSize, secondCache.Size(), "The size of the caches should match")
+	require.True(t, firstCache.Equal(secondCache), "The caches should be equal")
+}
+
 func BenchmarkAdd(b *testing.B) {
 
 	log.SetLogger("BenchmarkAdd", log.SILENT)

util/bytes.go

@@ -16,7 +16,9 @@
 
 package util
 
-import "encoding/binary"
+import (
+	"encoding/binary"
+)
 
 func Uint64AsBytes(i uint64) []byte {
 	b := make([]byte, 8)