119 lines
2.5 KiB
Go
119 lines
2.5 KiB
Go
package common
|
|
|
|
import (
|
|
"bytes"
|
|
"sync"
|
|
|
|
"github.com/golang/protobuf/proto"
|
|
"github.com/zcash-hackworks/lightwalletd/walletrpc"
|
|
)
|
|
|
|
type BlockCacheEntry struct {
|
|
data []byte
|
|
hash []byte
|
|
}
|
|
|
|
type BlockCache struct {
|
|
MaxEntries int
|
|
|
|
// m[firstBlock..nextBlock) are valid
|
|
m map[int]*BlockCacheEntry
|
|
firstBlock int
|
|
nextBlock int
|
|
|
|
mutex sync.RWMutex
|
|
}
|
|
|
|
func NewBlockCache(maxEntries int) *BlockCache {
|
|
return &BlockCache{
|
|
MaxEntries: maxEntries,
|
|
m: make(map[int]*BlockCacheEntry),
|
|
}
|
|
}
|
|
|
|
func (c *BlockCache) Add(height int, block *walletrpc.CompactBlock) (error, bool) {
|
|
// Invariant: m[firstBlock..nextBlock) are valid.
|
|
c.mutex.Lock()
|
|
defer c.mutex.Unlock()
|
|
|
|
if height > c.nextBlock {
|
|
// restarting the cache (never happens currently), or first time
|
|
for i := c.firstBlock; i < c.nextBlock; i++ {
|
|
delete(c.m, i)
|
|
}
|
|
c.firstBlock = height
|
|
c.nextBlock = height
|
|
}
|
|
// Invariant: m[firstBlock..nextBlock) are valid.
|
|
|
|
// If we already have this block, a reorg must have occurred;
|
|
// this block (and all higher) must be re-added.
|
|
h := height
|
|
if h < c.firstBlock {
|
|
h = c.firstBlock
|
|
}
|
|
for i := h; i < c.nextBlock; i++ {
|
|
delete(c.m, i)
|
|
}
|
|
c.nextBlock = height
|
|
if c.firstBlock > c.nextBlock {
|
|
c.firstBlock = c.nextBlock
|
|
}
|
|
// Invariant: m[firstBlock..nextBlock) are valid.
|
|
|
|
// Detect reorg, ingestor needs to handle it
|
|
if height > c.firstBlock && !bytes.Equal(block.PrevHash, c.m[height-1].hash) {
|
|
return nil, true
|
|
}
|
|
|
|
// Add the entry and update the counters
|
|
data, err := proto.Marshal(block)
|
|
if err != nil {
|
|
println("Error marshalling block!")
|
|
return err, false
|
|
}
|
|
c.m[height] = &BlockCacheEntry{
|
|
data: data,
|
|
hash: block.GetHash(),
|
|
}
|
|
c.nextBlock++
|
|
// Invariant: m[firstBlock..nextBlock) are valid.
|
|
|
|
// remove any blocks that are older than the capacity of the cache
|
|
for c.firstBlock < c.nextBlock-c.MaxEntries {
|
|
// Invariant: m[firstBlock..nextBlock) are valid.
|
|
delete(c.m, c.firstBlock)
|
|
c.firstBlock++
|
|
}
|
|
// Invariant: m[firstBlock..nextBlock) are valid.
|
|
|
|
return nil, false
|
|
}
|
|
|
|
func (c *BlockCache) Get(height int) *walletrpc.CompactBlock {
|
|
c.mutex.RLock()
|
|
defer c.mutex.RUnlock()
|
|
|
|
if height < c.firstBlock || height >= c.nextBlock {
|
|
return nil
|
|
}
|
|
|
|
serialized := &walletrpc.CompactBlock{}
|
|
err := proto.Unmarshal(c.m[height].data, serialized)
|
|
if err != nil {
|
|
println("Error unmarshalling compact block")
|
|
return nil
|
|
}
|
|
|
|
return serialized
|
|
}
|
|
|
|
func (c *BlockCache) GetLatestBlock() int {
|
|
c.mutex.RLock()
|
|
defer c.mutex.RUnlock()
|
|
if c.firstBlock == c.nextBlock {
|
|
return -1
|
|
}
|
|
return c.nextBlock - 1
|
|
}
|