Added caching for block chain. Currently set to 10k

4 files changed, 136 insertions, 2 deletions

diff --git a/core/block_cache.go b/core/block_cache.go
new file mode 100644
index 000000000..321021eb4
--- /dev/null
+++ b/core/block_cache.go
@@ -0,0 +1,68 @@
+package core
+
+import (
+	"sync"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/types"
+)
+
+// BlockCache implements a caching mechanism specifically for blocks and uses FILO to pop
+type BlockCache struct {
+	size int
+
+	hashes []common.Hash
+	blocks map[common.Hash]*types.Block
+
+	mu sync.RWMutex
+}
+
+// Creates and returns a `BlockCache` with `size`. If `size` is smaller than 1 it will panic
+func NewBlockCache(size int) *BlockCache {
+	if size < 1 {
+		panic("block cache size not allowed to be smaller than 1")
+	}
+
+	bc := &BlockCache{size: size}
+	bc.Clear()
+	return bc
+}
+
+func (bc *BlockCache) Clear() {
+	bc.blocks = make(map[common.Hash]*types.Block)
+	bc.hashes = nil
+
+}
+
+func (bc *BlockCache) Push(block *types.Block) {
+	bc.mu.Lock()
+	defer bc.mu.Unlock()
+
+	if len(bc.hashes) == bc.size {
+		delete(bc.blocks, bc.hashes[0])
+
+		// XXX There are a few other options on solving this
+		// 1) use a poller / GC like mechanism to clean up untracked objects
+		// 2) copy as below
+		// re-use the slice and remove the reference to bc.hashes[0]
+		// this will allow the element to be garbage collected.
+		copy(bc.hashes, bc.hashes[1:])
+	} else {
+		bc.hashes = append(bc.hashes, common.Hash{})
+	}
+
+	hash := block.Hash()
+	bc.blocks[hash] = block
+	bc.hashes[len(bc.hashes)-1] = hash
+}
+
+func (bc *BlockCache) Get(hash common.Hash) *types.Block {
+	bc.mu.RLock()
+	defer bc.mu.RUnlock()
+
+	if block, haz := bc.blocks[hash]; haz {
+		return block
+	}
+
+	return nil
+}
diff --git a/core/block_cache_test.go b/core/block_cache_test.go
new file mode 100644
index 000000000..d4f610b71
--- /dev/null
+++ b/core/block_cache_test.go
@@ -0,0 +1,48 @@
+package core
+
+import (
+	"math/big"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/types"
+)
+
+func newChain(size int) (chain []*types.Block) {
+	var parentHash common.Hash
+	for i := 0; i < size; i++ {
+		block := types.NewBlock(parentHash, common.Address{}, common.Hash{}, new(big.Int), 0, "")
+		block.Header().Number = big.NewInt(int64(i))
+		chain = append(chain, block)
+		parentHash = block.Hash()
+	}
+	return
+}
+
+func insertChainCache(cache *BlockCache, chain []*types.Block) {
+	for _, block := range chain {
+		cache.Push(block)
+	}
+}
+
+func TestNewBlockCache(t *testing.T) {
+	chain := newChain(3)
+	cache := NewBlockCache(2)
+	insertChainCache(cache, chain)
+
+	if cache.hashes[0] != chain[1].Hash() {
+		t.Error("oldest block incorrect")
+	}
+}
+
+func TestInclusion(t *testing.T) {
+	chain := newChain(3)
+	cache := NewBlockCache(3)
+	insertChainCache(cache, chain)
+
+	for _, block := range chain {
+		if b := cache.Get(block.Hash()); b == nil {
+			t.Errorf("getting %x failed", block.Hash())
+		}
+	}
+}
diff --git a/core/chain_manager.go b/core/chain_manager.go
index 915fa704f..1bc8edea6 100644
--- a/core/chain_manager.go
+++ b/core/chain_manager.go
@@ -23,6 +23,8 @@ var (
 	blockNumPre  = []byte("block-num-")
 )
 
+const blockCacheLimit = 10000
+
 type StateQuery interface {
 	GetAccount(addr []byte) *state.StateObject
 }
@@ -92,15 +94,25 @@ type ChainManager struct {
 	transState *state.StateDB
 	txState    *state.ManagedState
 
+	cache *BlockCache
+
 	quit chan struct{}
 }
 
 func NewChainManager(blockDb, stateDb common.Database, mux *event.TypeMux) *ChainManager {
-	bc := &ChainManager{blockDb: blockDb, stateDb: stateDb, genesisBlock: GenesisBlock(stateDb), eventMux: mux, quit: make(chan struct{})}
+	bc := &ChainManager{blockDb: blockDb, stateDb: stateDb, genesisBlock: GenesisBlock(stateDb), eventMux: mux, quit: make(chan struct{}), cache: NewBlockCache(blockCacheLimit)}
 	bc.setLastBlock()
 	bc.transState = bc.State().Copy()
 	// Take ownership of this particular state
 	bc.txState = state.ManageState(bc.State().Copy())
+
+	// load in last `blockCacheLimit` - 1 blocks. Last block is the current.
+	ancestors := bc.GetAncestors(bc.currentBlock, blockCacheLimit-1)
+	ancestors = append(ancestors, bc.currentBlock)
+	for _, block := range ancestors {
+		bc.cache.Push(block)
+	}
+
 	go bc.update()
 
 	return bc
@@ -275,6 +287,8 @@ func (bc *ChainManager) insert(block *types.Block) {
 
 	key := append(blockNumPre, block.Number().Bytes()...)
 	bc.blockDb.Put(key, bc.lastBlockHash.Bytes())
+	// Push block to cache
+	bc.cache.Push(block)
 }
 
 func (bc *ChainManager) write(block *types.Block) {
@@ -318,6 +332,10 @@ func (self *ChainManager) GetBlockHashesFromHash(hash common.Hash, max uint64) (
 }
 
 func (self *ChainManager) GetBlock(hash common.Hash) *types.Block {
+	if block := self.cache.Get(hash); block != nil {
+		return block
+	}
+
 	data, _ := self.blockDb.Get(append(blockHashPre, hash[:]...))
 	if len(data) == 0 {
 		return nil
diff --git a/core/chain_manager_test.go b/core/chain_manager_test.go
index e49e594a3..bf172f3bf 100644
--- a/core/chain_manager_test.go
+++ b/core/chain_manager_test.go
@@ -69,7 +69,7 @@ func printChain(bc *ChainManager) {
 func testChain(chainB types.Blocks, bman *BlockProcessor) (*big.Int, error) {
 	td := new(big.Int)
 	for _, block := range chainB {
-		td2, err := bman.bc.processor.Process(block)
+		td2, _, err := bman.bc.processor.Process(block)
 		if err != nil {
 			if IsKnownBlockErr(err) {
 				continue