lite2: prune-headers (#4478)

closes #4469

Improved speed of cleanup by using SignedHeaderAfter instead of TrustedHeader to jump from header to header.

Prune() is now called when a new header and validator set are saved and is a function dealt by the database itself

## Commits:

* prune headers and vals

* modified cleanup and tests

* fixes after my own review

* implement Prune func

* make db ops concurrently safe

* use Iterator in SignedHeaderAfter

we should iterate from height+1, not from the end!

* simplify cleanup

Co-authored-by: Anton Kaliaev <anton.kalyaev@gmail.com>

lite2/client.go

@@ -22,9 +22,9 @@ const (
 	sequential mode = iota + 1
 	skipping
 
-	defaultUpdatePeriod                       = 5 * time.Second
-	defaultRemoveNoLongerTrustedHeadersPeriod = 24 * time.Hour
-	defaultMaxRetryAttempts                   = 10
+	defaultUpdatePeriod     = 5 * time.Second
+	defaultPruningSize      = 1000
+	defaultMaxRetryAttempts = 10
 )
 
 // Option sets a parameter for the light client.
@@ -62,13 +62,14 @@ func UpdatePeriod(d time.Duration) Option {
 	}
 }
 
-// RemoveNoLongerTrustedHeadersPeriod option can be used to define how often
-// the routine, which cleans up no longer trusted headers (outside of trusting
-// period), is run. Default: once a day. When set to zero, the routine won't be
-// started.
-func RemoveNoLongerTrustedHeadersPeriod(d time.Duration) Option {
+// PruningSize option sets the maximum amount of headers & validator set pairs
+// that the light client stores. When Prune() is run, all headers (along with
+// the associated validator sets) that are earlier than the h amount of headers
+// will be removed from the store. Default: 1000. A pruning size of 0 will not
+// prune the lite client at all.
+func PruningSize(h uint16) Option {
 	return func(c *Client) {
-		c.removeNoLongerTrustedHeadersPeriod = d
+		c.pruningSize = h
 	}
 }
 
@@ -128,7 +129,7 @@ type Client struct {
 	// See UpdatePeriod option
 	updatePeriod time.Duration
 	// See RemoveNoLongerTrustedHeadersPeriod option
-	removeNoLongerTrustedHeadersPeriod time.Duration
+	pruningSize uint16
 	// See ConfirmationFunction option
 	confirmationFn func(action string) bool
 
@@ -191,19 +192,19 @@ func NewClientFromTrustedStore(
 	options ...Option) (*Client, error) {
 
 	c := &Client{
-		chainID:                            chainID,
-		trustingPeriod:                     trustingPeriod,
-		verificationMode:                   skipping,
-		trustLevel:                         DefaultTrustLevel,
-		maxRetryAttempts:                   defaultMaxRetryAttempts,
-		primary:                            primary,
-		witnesses:                          witnesses,
-		trustedStore:                       trustedStore,
-		updatePeriod:                       defaultUpdatePeriod,
-		removeNoLongerTrustedHeadersPeriod: defaultRemoveNoLongerTrustedHeadersPeriod,
-		confirmationFn:                     func(action string) bool { return true },
-		quit:                               make(chan struct{}),
-		logger:                             log.NewNopLogger(),
+		chainID:          chainID,
+		trustingPeriod:   trustingPeriod,
+		verificationMode: skipping,
+		trustLevel:       DefaultTrustLevel,
+		maxRetryAttempts: defaultMaxRetryAttempts,
+		primary:          primary,
+		witnesses:        witnesses,
+		trustedStore:     trustedStore,
+		updatePeriod:     defaultUpdatePeriod,
+		pruningSize:      defaultPruningSize,
+		confirmationFn:   func(action string) bool { return true },
+		quit:             make(chan struct{}),
+		logger:           log.NewNopLogger(),
 	}
 
 	for _, o := range options {
@@ -304,7 +305,10 @@ func (c *Client) checkTrustedHeaderUsingOptions(options TrustOptions) error {
 			c.latestTrustedHeader.Height, c.latestTrustedHeader.Hash())
 		if c.confirmationFn(action) {
 			// remove all the headers (options.Height, trustedHeader.Height]
-			c.cleanup(options.Height + 1)
+			err := c.cleanupAfter(options.Height)
+			if err != nil {
+				return errors.Wrapf(err, "cleanupAfter(%d)", options.Height)
+			}
 
 			c.logger.Info("Rolled back to older header (newer headers were removed)",
 				"old", options.Height)
@@ -388,11 +392,6 @@ func (c *Client) initializeWithTrustOptions(options TrustOptions) error {
 func (c *Client) Start() error {
 	c.logger.Info("Starting light client")
 
-	if c.removeNoLongerTrustedHeadersPeriod > 0 {
-		c.routinesWaitGroup.Add(1)
-		go c.removeNoLongerTrustedHeadersRoutine()
-	}
-
 	if c.updatePeriod > 0 {
 		c.routinesWaitGroup.Add(1)
 		go c.autoUpdateRoutine()
@@ -557,6 +556,7 @@ func (c *Client) VerifyHeader(newHeader *types.SignedHeader, newVals *types.Vali
 	// Check if newHeader already verified.
 	h, err := c.TrustedHeader(newHeader.Height)
 	if err == nil {
+		// Make sure it's the same header.
 		if !bytes.Equal(h.Hash(), newHeader.Hash()) {
 			return errors.Errorf("existing trusted header %X does not match newHeader %X", h.Hash(), newHeader.Hash())
 		}
@@ -630,40 +630,36 @@ func (c *Client) Witnesses() []provider.Provider {
 // client must be stopped at this point.
 func (c *Client) Cleanup() error {
 	c.logger.Info("Removing all the data")
-	return c.cleanup(0)
+	c.latestTrustedHeader = nil
+	c.latestTrustedVals = nil
+	return c.trustedStore.Prune(0)
 }
 
-// cleanup deletes all headers & validator sets between +stopHeight+ and latest
-// height included. It also sets trustedHeader (vals) to the latest header
-// (vals) if such exists.
-func (c *Client) cleanup(stopHeight int64) error {
-	// 1) Get the oldest height.
-	oldestHeight, err := c.trustedStore.FirstSignedHeaderHeight()
-	if err != nil {
-		return errors.Wrap(err, "can't get first trusted height")
-	}
+// cleanupAfter deletes all headers & validator sets after +height+. It also
+// resets latestTrustedHeader to the latest header.
+func (c *Client) cleanupAfter(height int64) error {
+	nextHeight := height
 
-	// 2) Get the latest height.
-	latestHeight, err := c.trustedStore.LastSignedHeaderHeight()
-	if err != nil {
-		return errors.Wrap(err, "can't get last trusted height")
-	}
-
-	// 3) Remove all headers and validator sets.
-	if stopHeight < oldestHeight {
-		stopHeight = oldestHeight
-	}
-	for height := stopHeight; height <= latestHeight; height++ {
-		err = c.trustedStore.DeleteSignedHeaderAndValidatorSet(height)
-		if err != nil {
-			c.logger.Error("can't remove a trusted header & validator set", "err", err, "height", height)
-			continue
+	for {
+		h, err := c.trustedStore.SignedHeaderAfter(nextHeight)
+		if err == store.ErrSignedHeaderNotFound {
+			break
+		} else if err != nil {
+			return errors.Wrapf(err, "failed to get header after %d", nextHeight)
 		}
+
+		err = c.trustedStore.DeleteSignedHeaderAndValidatorSet(h.Height)
+		if err != nil {
+			c.logger.Error("can't remove a trusted header & validator set", "err", err,
+				"height", h.Height)
+		}
+
+		nextHeight = h.Height
 	}
 
 	c.latestTrustedHeader = nil
 	c.latestTrustedVals = nil
-	err = c.restoreTrustedHeaderAndVals()
+	err := c.restoreTrustedHeaderAndVals()
 	if err != nil {
 		return err
 	}
@@ -776,7 +772,12 @@ func (c *Client) updateTrustedHeaderAndVals(h *types.SignedHeader, vals *types.V
 		return errors.Wrap(err, "failed to save trusted header")
 	}
 
-	// Only update latestTrustedHeader if we move forward (not backwards).
+	if c.pruningSize > 0 {
+		if err := c.trustedStore.Prune(c.pruningSize); err != nil {
+			return errors.Wrap(err, "prune")
+		}
+	}
+
 	if c.latestTrustedHeader == nil || h.Height > c.latestTrustedHeader.Height {
 		c.latestTrustedHeader = h
 		c.latestTrustedVals = vals
@@ -921,72 +922,6 @@ func (c *Client) removeWitness(idx int) {
 	}
 }
 
-func (c *Client) removeNoLongerTrustedHeadersRoutine() {
-	defer c.routinesWaitGroup.Done()
-
-	ticker := time.NewTicker(c.removeNoLongerTrustedHeadersPeriod)
-	defer ticker.Stop()
-
-	for {
-		select {
-		case <-ticker.C:
-			c.RemoveNoLongerTrustedHeaders(time.Now())
-		case <-c.quit:
-			return
-		}
-	}
-}
-
-// RemoveNoLongerTrustedHeaders removes no longer trusted headers (due to
-// expiration).
-//
-// Exposed for testing.
-func (c *Client) RemoveNoLongerTrustedHeaders(now time.Time) {
-	// 1) Get the oldest height.
-	oldestHeight, err := c.FirstTrustedHeight()
-	if err != nil {
-		c.logger.Error("can't get first trusted height", "err", err)
-		return
-	}
-	if oldestHeight == -1 { // no headers yet => wait
-		return
-	}
-
-	// 2) Get the latest height.
-	latestHeight, err := c.LastTrustedHeight()
-	if err != nil {
-		c.logger.Error("can't get last trusted height", "err", err)
-		return
-	}
-	if latestHeight == -1 { // no headers yet => wait
-		return
-	}
-
-	// 3) Remove all headers that are outside of the trusting period.
-	//
-	// NOTE: even the latest header can be removed. it's okay because
-	// c.latestTrustedHeader will retain it in memory so other funcs like VerifyHeader
-	// don't crash.
-	for height := oldestHeight; height <= latestHeight; height++ {
-		h, err := c.trustedStore.SignedHeader(height)
-		if err != nil {
-			c.logger.Error("can't get a trusted header", "err", err, "height", height)
-			continue
-		}
-
-		// Stop if the header is within the trusting period.
-		if !HeaderExpired(h, c.trustingPeriod, now) {
-			break
-		}
-
-		err = c.trustedStore.DeleteSignedHeaderAndValidatorSet(height)
-		if err != nil {
-			c.logger.Error("can't remove a trusted header & validator set", "err", err, "height", height)
-			continue
-		}
-	}
-}
-
 func (c *Client) autoUpdateRoutine() {
 	defer c.routinesWaitGroup.Done()
 

lite2/client_test.go

@@ -306,47 +306,6 @@ func TestClient_SkippingVerification(t *testing.T) {
 	}
 }
 
-func TestClientRemovesNoLongerTrustedHeaders(t *testing.T) {
-	c, err := NewClient(
-		chainID,
-		trustOptions,
-		fullNode,
-		[]provider.Provider{fullNode},
-		dbs.New(dbm.NewMemDB(), chainID),
-		Logger(log.TestingLogger()),
-	)
-
-	assert.NotPanics(t, func() {
-		now := bTime.Add(4 * time.Hour).Add(1 * time.Second)
-		c.RemoveNoLongerTrustedHeaders(now)
-	})
-
-	require.NoError(t, err)
-	err = c.Start()
-	require.NoError(t, err)
-	defer c.Stop()
-
-	// Verify new headers.
-	_, err = c.VerifyHeaderAtHeight(2, bTime.Add(2*time.Hour).Add(1*time.Second))
-	require.NoError(t, err)
-	now := bTime.Add(4 * time.Hour).Add(1 * time.Second)
-	_, err = c.VerifyHeaderAtHeight(3, now)
-	require.NoError(t, err)
-
-	// Remove expired headers.
-	c.RemoveNoLongerTrustedHeaders(now)
-
-	// Check expired headers are no longer available.
-	h, err := c.TrustedHeader(1)
-	assert.Error(t, err)
-	assert.Nil(t, h)
-
-	// Check not expired headers are available.
-	h, err = c.TrustedHeader(2)
-	assert.NoError(t, err)
-	assert.NotNil(t, h)
-}
-
 func TestClient_Cleanup(t *testing.T) {
 	c, err := NewClient(
 		chainID,
@@ -357,10 +316,9 @@ func TestClient_Cleanup(t *testing.T) {
 		Logger(log.TestingLogger()),
 	)
 	require.NoError(t, err)
-	err = c.Start()
+	_, err = c.TrustedHeader(1)
 	require.NoError(t, err)
 
-	c.Stop()
 	err = c.Cleanup()
 	require.NoError(t, err)
 
@@ -387,9 +345,6 @@ func TestClientRestoresTrustedHeaderAfterStartup1(t *testing.T) {
 			Logger(log.TestingLogger()),
 		)
 		require.NoError(t, err)
-		err = c.Start()
-		require.NoError(t, err)
-		defer c.Stop()
 
 		h, err := c.TrustedHeader(1)
 		assert.NoError(t, err)
@@ -429,14 +384,12 @@ func TestClientRestoresTrustedHeaderAfterStartup1(t *testing.T) {
 			Logger(log.TestingLogger()),
 		)
 		require.NoError(t, err)
-		err = c.Start()
-		require.NoError(t, err)
-		defer c.Stop()
 
 		h, err := c.TrustedHeader(1)
 		assert.NoError(t, err)
-		assert.NotNil(t, h)
-		assert.Equal(t, h.Hash(), header1.Hash())
+		if assert.NotNil(t, h) {
+			assert.Equal(t, h.Hash(), header1.Hash())
+		}
 	}
 }
 

lite2/store/db/db.go

@@ -1,9 +1,11 @@
 package db
 
 import (
+	"encoding/binary"
 	"fmt"
 	"regexp"
 	"strconv"
+	"sync"
 
 	"github.com/pkg/errors"
 	"github.com/tendermint/go-amino"
@@ -14,10 +16,17 @@ import (
 	"github.com/tendermint/tendermint/types"
 )
 
+var (
+	sizeKey = []byte("size")
+)
+
 type dbs struct {
 	db     dbm.DB
 	prefix string
 
+	mtx  sync.RWMutex
+	size uint16
+
 	cdc *amino.Codec
 }
 
@@ -28,11 +37,20 @@ type dbs struct {
 func New(db dbm.DB, prefix string) store.Store {
 	cdc := amino.NewCodec()
 	cryptoAmino.RegisterAmino(cdc)
-	return &dbs{db: db, prefix: prefix, cdc: cdc}
+
+	size := uint16(0)
+	bz, err := db.Get(sizeKey)
+	if err == nil && len(bz) > 0 {
+		size = unmarshalSize(bz)
+	}
+
+	return &dbs{db: db, prefix: prefix, cdc: cdc, size: size}
 }
 
 // SaveSignedHeaderAndValidatorSet persists SignedHeader and ValidatorSet to
 // the db.
+//
+// Safe for concurrent use by multiple goroutines.
 func (s *dbs) SaveSignedHeaderAndValidatorSet(sh *types.SignedHeader, valSet *types.ValidatorSet) error {
 	if sh.Height <= 0 {
 		panic("negative or zero height")
@@ -48,30 +66,54 @@ func (s *dbs) SaveSignedHeaderAndValidatorSet(sh *types.SignedHeader, valSet *ty
 		return errors.Wrap(err, "marshalling validator set")
 	}
 
+	s.mtx.Lock()
+	defer s.mtx.Unlock()
+
 	b := s.db.NewBatch()
 	b.Set(s.shKey(sh.Height), shBz)
 	b.Set(s.vsKey(sh.Height), valSetBz)
+	b.Set(sizeKey, marshalSize(s.size+1))
+
 	err = b.WriteSync()
 	b.Close()
+
+	if err == nil {
+		s.size++
+	}
+
 	return err
 }
 
 // DeleteSignedHeaderAndValidatorSet deletes SignedHeader and ValidatorSet from
 // the db.
+//
+// Safe for concurrent use by multiple goroutines.
 func (s *dbs) DeleteSignedHeaderAndValidatorSet(height int64) error {
 	if height <= 0 {
 		panic("negative or zero height")
 	}
 
+	s.mtx.Lock()
+	defer s.mtx.Unlock()
+
 	b := s.db.NewBatch()
 	b.Delete(s.shKey(height))
 	b.Delete(s.vsKey(height))
+	b.Set(sizeKey, marshalSize(s.size-1))
+
 	err := b.WriteSync()
 	b.Close()
+
+	if err == nil {
+		s.size--
+	}
+
 	return err
 }
 
 // SignedHeader loads SignedHeader at the given height.
+//
+// Safe for concurrent use by multiple goroutines.
 func (s *dbs) SignedHeader(height int64) (*types.SignedHeader, error) {
 	if height <= 0 {
 		panic("negative or zero height")
@@ -91,6 +133,8 @@ func (s *dbs) SignedHeader(height int64) (*types.SignedHeader, error) {
 }
 
 // ValidatorSet loads ValidatorSet at the given height.
+//
+// Safe for concurrent use by multiple goroutines.
 func (s *dbs) ValidatorSet(height int64) (*types.ValidatorSet, error) {
 	if height <= 0 {
 		panic("negative or zero height")
@@ -110,6 +154,8 @@ func (s *dbs) ValidatorSet(height int64) (*types.ValidatorSet, error) {
 }
 
 // LastSignedHeaderHeight returns the last SignedHeader height stored.
+//
+// Safe for concurrent use by multiple goroutines.
 func (s *dbs) LastSignedHeaderHeight() (int64, error) {
 	itr, err := s.db.ReverseIterator(
 		s.shKey(1),
@@ -133,6 +179,8 @@ func (s *dbs) LastSignedHeaderHeight() (int64, error) {
 }
 
 // FirstSignedHeaderHeight returns the first SignedHeader height stored.
+//
+// Safe for concurrent use by multiple goroutines.
 func (s *dbs) FirstSignedHeaderHeight() (int64, error) {
 	itr, err := s.db.Iterator(
 		s.shKey(1),
@@ -155,12 +203,16 @@ func (s *dbs) FirstSignedHeaderHeight() (int64, error) {
 	return -1, nil
 }
 
+// SignedHeaderAfter iterates over headers until it finds a header after one at
+// height. It returns ErrSignedHeaderNotFound if no such header exists.
+//
+// Safe for concurrent use by multiple goroutines.
 func (s *dbs) SignedHeaderAfter(height int64) (*types.SignedHeader, error) {
 	if height <= 0 {
 		panic("negative or zero height")
 	}
 
-	itr, err := s.db.ReverseIterator(
+	itr, err := s.db.Iterator(
 		s.shKey(height+1),
 		append(s.shKey(1<<63-1), byte(0x00)),
 	)
@@ -178,7 +230,76 @@ func (s *dbs) SignedHeaderAfter(height int64) (*types.SignedHeader, error) {
 		itr.Next()
 	}
 
-	panic(fmt.Sprintf("no header after height %d. make sure height is not greater than latest existing height", height))
+	return nil, store.ErrSignedHeaderNotFound
+}
+
+// Prune prunes header & validator set pairs until there are only size pairs
+// left.
+//
+// Safe for concurrent use by multiple goroutines.
+func (s *dbs) Prune(size uint16) error {
+	// 1) Check how many we need to prune.
+	s.mtx.RLock()
+	sSize := s.size
+	s.mtx.RUnlock()
+
+	if sSize <= size { // nothing to prune
+		return nil
+	}
+	numToPrune := sSize - size
+
+	// 2) Iterate over headers and perform a batch operation.
+	itr, err := s.db.Iterator(
+		s.shKey(1),
+		append(s.shKey(1<<63-1), byte(0x00)),
+	)
+	if err != nil {
+		panic(err)
+	}
+
+	b := s.db.NewBatch()
+
+	pruned := 0
+	for itr.Valid() && numToPrune > 0 {
+		key := itr.Key()
+		_, height, ok := parseShKey(key)
+		if ok {
+			b.Delete(s.shKey(height))
+			b.Delete(s.vsKey(height))
+		}
+		itr.Next()
+		numToPrune--
+		pruned++
+	}
+
+	itr.Close()
+
+	err = b.WriteSync()
+	b.Close()
+	if err != nil {
+		return err
+	}
+
+	// 3) Update size.
+	s.mtx.Lock()
+	defer s.mtx.Unlock()
+
+	s.size -= uint16(pruned)
+
+	if wErr := s.db.SetSync(sizeKey, marshalSize(s.size)); wErr != nil {
+		return errors.Wrap(wErr, "failed to persist size")
+	}
+
+	return nil
+}
+
+// Size returns the number of header & validator set pairs.
+//
+// Safe for concurrent use by multiple goroutines.
+func (s *dbs) Size() uint16 {
+	s.mtx.RLock()
+	defer s.mtx.RUnlock()
+	return s.size
 }
 
 func (s *dbs) shKey(height int64) []byte {
@@ -214,3 +335,13 @@ func parseShKey(key []byte) (prefix string, height int64, ok bool) {
 	}
 	return
 }
+
+func marshalSize(size uint16) []byte {
+	bs := make([]byte, 2)
+	binary.LittleEndian.PutUint16(bs, size)
+	return bs
+}
+
+func unmarshalSize(bz []byte) uint16 {
+	return binary.LittleEndian.Uint16(bz)
+}

lite2/store/db/db_test.go

@@ -1,6 +1,7 @@
 package db
 
 import (
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -93,3 +94,69 @@ func Test_SignedHeaderAfter(t *testing.T) {
 		assert.EqualValues(t, 2, h.Height)
 	}
 }
+
+func Test_Prune(t *testing.T) {
+	dbStore := New(dbm.NewMemDB(), "Test_Prune")
+
+	// Empty store
+	assert.EqualValues(t, 0, dbStore.Size())
+	err := dbStore.Prune(0)
+	require.NoError(t, err)
+
+	// One header
+	err = dbStore.SaveSignedHeaderAndValidatorSet(
+		&types.SignedHeader{Header: &types.Header{Height: 2}}, &types.ValidatorSet{})
+	require.NoError(t, err)
+
+	assert.EqualValues(t, 1, dbStore.Size())
+
+	err = dbStore.Prune(1)
+	require.NoError(t, err)
+	assert.EqualValues(t, 1, dbStore.Size())
+
+	err = dbStore.Prune(0)
+	require.NoError(t, err)
+	assert.EqualValues(t, 0, dbStore.Size())
+
+	// Multiple headers
+	for i := 1; i <= 10; i++ {
+		err = dbStore.SaveSignedHeaderAndValidatorSet(
+			&types.SignedHeader{Header: &types.Header{Height: int64(i)}}, &types.ValidatorSet{})
+		require.NoError(t, err)
+	}
+
+	err = dbStore.Prune(11)
+	require.NoError(t, err)
+	assert.EqualValues(t, 10, dbStore.Size())
+
+	err = dbStore.Prune(7)
+	require.NoError(t, err)
+	assert.EqualValues(t, 7, dbStore.Size())
+}
+
+func Test_Concurrency(t *testing.T) {
+	dbStore := New(dbm.NewMemDB(), "Test_Prune")
+
+	var wg sync.WaitGroup
+	for i := 1; i <= 100; i++ {
+		wg.Add(1)
+		go func(i int64) {
+			defer wg.Done()
+
+			dbStore.SaveSignedHeaderAndValidatorSet(
+				&types.SignedHeader{Header: &types.Header{Height: i}}, &types.ValidatorSet{})
+
+			dbStore.SignedHeader(i)
+			dbStore.ValidatorSet(i)
+			dbStore.LastSignedHeaderHeight()
+			dbStore.FirstSignedHeaderHeight()
+
+			dbStore.Prune(2)
+			_ = dbStore.Size()
+
+			dbStore.DeleteSignedHeaderAndValidatorSet(1)
+		}(int64(i))
+	}
+
+	wg.Wait()
+}

lite2/store/store.go

@@ -45,4 +45,11 @@ type Store interface {
 	//
 	// height must be > 0 && <= LastSignedHeaderHeight.
 	SignedHeaderAfter(height int64) (*types.SignedHeader, error)
+
+	// Prune removes headers & the associated validator sets when Store reaches a
+	// defined size (number of header & validator set pairs).
+	Prune(size uint16) error
+
+	// Size returns a number of currently existing header & validator set pairs.
+	Size() uint16
 }