lightwalletd/common/cache.go

package common

import (
	"bytes"
	"sync"
	"time"

	"git.hush.is/hush/lightwalletd/walletrpc"
	"github.com/golang/protobuf/proto"
)

type BlockCacheEntry struct {
	data []byte
	hash []byte
}

type BlockCache struct {
	MaxEntries int

	FirstBlock int
	LastBlock  int

	m map[int]*BlockCacheEntry

	mutex sync.RWMutex

	// Adaptive confirmation lag. lightwalletd advertises a tip that trails the
	// real tip by a number of blocks that scales with the recent reorg rate, so
	// wallets anchor shielded spends at a settled height while the chain is
	// churning, and at (near) the real tip when it is stable. reorgTimes holds
	// the times of recent reorg detections, pruned to lagWindow.
	reorgTimes  []time.Time
	lagWindow   time.Duration
	lagMin      int
	lagMax      int
	adaptiveLag bool

	// advTip is the cached advertised (lag-adjusted) tip, kept fresh by Add so
	// the per-block serving cap (HeightAllowed) is a cheap read. -1 until ready.
	advTip int
}

func NewBlockCache(maxEntries int) *BlockCache {
	return &BlockCache{
		MaxEntries:  maxEntries,
		FirstBlock:  -1,
		LastBlock:   -1,
		m:           make(map[int]*BlockCacheEntry),
		adaptiveLag: true,
		lagWindow:   30 * time.Minute,
		lagMin:      1,
		lagMax:      12,
		advTip:      -1,
	}
}

func (c *BlockCache) Add(height int, block *walletrpc.CompactBlock) (error, bool) {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	//println("Cache add", height)
	if c.FirstBlock == -1 && c.LastBlock == -1 {
		// If this is the first block, prep the data structure
		c.FirstBlock = height
		c.LastBlock = height - 1
	}

	// If we're adding a block in the middle of the cache, remove all
	// blocks after it, since this might be a reorg, and we don't want
	// Any outdated blocks returned
	if height >= c.FirstBlock && height <= c.LastBlock {
		for i := height; i <= c.LastBlock; i++ {
			delete(c.m, i)
		}
		c.LastBlock = height - 1
	}

	// Don't allow out-of-order blocks. This is more of a sanity check than anything
	// If there is a reorg, then the ingestor needs to handle it.
	if c.m[height-1] != nil && !bytes.Equal(block.PrevHash, c.m[height-1].hash) {
		// Record the reorg so the adaptive confirmation lag can react to it.
		c.reorgTimes = append(c.reorgTimes, time.Now())
		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.LastBlock = height

	// If the cache is full, remove the oldest block
	if c.LastBlock-c.FirstBlock+1 > c.MaxEntries {
		//println("Deleteing at height", c.FirstBlock)
		delete(c.m, c.FirstBlock)
		c.FirstBlock = c.FirstBlock + 1
	}

	// Keep the advertised (lag-adjusted) tip fresh for the block-serving cap.
	c.advTip = c.LastBlock - c.computeLagLocked()
	if c.advTip < c.FirstBlock {
		c.advTip = c.FirstBlock
	}

	//println("Cache size is ", len(c.m))
	return nil, false
}

func (c *BlockCache) Get(height int) *walletrpc.CompactBlock {
	c.mutex.RLock()
	defer c.mutex.RUnlock()

	//println("Cache get", height)
	if c.LastBlock == -1 || c.FirstBlock == -1 {
		return nil
	}

	if height < c.FirstBlock || height > c.LastBlock {
		//println("Cache miss: index out of range")
		return nil
	}

	//println("Cache returned")
	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()

	return c.LastBlock
}

// ConfigureLag sets the adaptive-confirmation-lag parameters. Called once at
// startup from the command-line flags.
func (c *BlockCache) ConfigureLag(adaptive bool, windowMinutes, min, max int) {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	c.adaptiveLag = adaptive
	c.lagWindow = time.Duration(windowMinutes) * time.Minute
	c.lagMin = min
	c.lagMax = max
}

// computeLagLocked returns the confirmation lag (in blocks) to apply right now
// and prunes reorg events that have aged out of the window. The lag is the
// configured floor plus the number of reorgs seen within lagWindow, capped at
// lagMax. Caller must hold c.mutex.
func (c *BlockCache) computeLagLocked() int {
	if !c.adaptiveLag {
		return c.lagMin
	}
	cutoff := time.Now().Add(-c.lagWindow)
	kept := c.reorgTimes[:0]
	for _, t := range c.reorgTimes {
		if t.After(cutoff) {
			kept = append(kept, t)
		}
	}
	c.reorgTimes = kept
	lag := c.lagMin + len(kept)
	if lag > c.lagMax {
		lag = c.lagMax
	}
	if lag < c.lagMin {
		lag = c.lagMin
	}
	return lag
}

// CurrentLag returns the confirmation lag currently being applied.
func (c *BlockCache) CurrentLag() int {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	return c.computeLagLocked()
}

// AdvertisedLatestBlock is the tip height lightwalletd reports to wallets: the
// real cache tip minus the adaptive confirmation lag. Wallets sync to and
// anchor shielded spends at this settled height, immune to tip reorgs. The
// internal cache and ingestor continue to track the real tip.
func (c *BlockCache) AdvertisedLatestBlock() int {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	if c.LastBlock < 0 {
		return c.LastBlock
	}
	adv := c.LastBlock - c.computeLagLocked()
	if adv < c.FirstBlock {
		adv = c.FirstBlock
	}
	c.advTip = adv
	return adv
}

// HeightAllowed reports whether height is at or below the advertised
// (lag-adjusted) tip. lightwalletd refuses to serve blocks above it so wallets
// cannot sync or anchor shielded spends into the unstable reorg zone. Cheap
// (read lock, no recompute); advTip is kept fresh by Add.
func (c *BlockCache) HeightAllowed(height int) bool {
	c.mutex.RLock()
	defer c.mutex.RUnlock()
	if c.advTip < 0 {
		return true
	}
	return height <= c.advTip
}