Efficient Consensus State Iteration (#125)

* start with efficient consensus state lookup

* writeup pruning logic

* fix tests

* add documentation

* improve byte efficiency

* actually fix tests and bug

* deduplicate

* fix return

* edit changelog

CHANGELOG.md

@@ -57,6 +57,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 ### Improvements
 
 * (modules/core/04-channel) [\#7949](https://github.com/cosmos/cosmos-sdk/issues/7949) Standardized channel `Acknowledgement` moved to its own file. Codec registration redundancy removed.
+* (modules/light-clients/07-tendermint) [\#125](https://github.com/cosmos/ibc-go/pull/125) Implement efficient iteration of consensus states and pruning of earliest expired consensus state on UpdateClient.
 
 ## IBC in the Cosmos SDK Repository
 

modules/light-clients/07-tendermint/types/store.go

@@ -1,18 +1,42 @@
 package types
 
 import (
+	"encoding/binary"
 	"strings"
 
 	"github.com/cosmos/cosmos-sdk/codec"
+	"github.com/cosmos/cosmos-sdk/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	clienttypes "github.com/cosmos/ibc-go/modules/core/02-client/types"
 	host "github.com/cosmos/ibc-go/modules/core/24-host"
 	"github.com/cosmos/ibc-go/modules/core/exported"
 )
 
-// KeyProcessedTime is appended to consensus state key to store the processed time
-var KeyProcessedTime = []byte("/processedTime")
+/*
+This file contains the logic for storage and iteration over `IterationKey` metadata that is stored
+for each consensus state. The consensus state key specified in ICS-24 and expected by counterparty chains
+stores the consensus state under the key: `consensusStates/{revision_number}-{revision_height}`, with each number
+represented as a string.
+While this works fine for IBC proof verification, it makes efficient iteration difficult since the lexicographic order
+of the consensus state keys do not match the height order of consensus states. This makes consensus state pruning and
+monotonic time enforcement difficult since it is inefficient to find the earliest consensus state or to find the neigboring
+consensus states given a consensus state height.
+Changing the ICS-24 representation will be a major breaking change that requires counterparty chains to accept a new key format.
+Thus to avoid breaking IBC, we can store a lookup from a more efficiently formatted key: `iterationKey` to the consensus state key which
+stores the underlying consensus state. This efficient iteration key will be formatted like so: `iterateConsensusStates{BigEndianRevisionBytes}{BigEndianHeightBytes}`.
+This ensures that the lexicographic order of iteration keys match the height order of the consensus states. Thus, we can use the SDK store's
+Iterators to iterate over the consensus states in ascending/descending order by providing a mapping from `iterationKey -> consensusStateKey -> ConsensusState`.
+A future version of IBC may choose to replace the ICS24 ConsensusState path with the more efficient format and make this indirection unnecessary.
+*/
+
+const KeyIterateConsensusStatePrefix = "iterateConsensusStates"
+
+var (
+	// KeyProcessedTime is appended to consensus state key to store the processed time
+	KeyProcessedTime = []byte("/processedTime")
+	KeyIteration     = []byte("/iterationKey")
+)
 
 // SetConsensusState stores the consensus state at the given height.
 func SetConsensusState(clientStore sdk.KVStore, cdc codec.BinaryMarshaler, consensusState *ConsensusState, height exported.Height) {
@@ -48,6 +72,12 @@ func GetConsensusState(store sdk.KVStore, cdc codec.BinaryMarshaler, height expo
 	return consensusState, nil
 }
 
+// deleteConsensusState deletes the consensus state at the given height
+func deleteConsensusState(clientStore sdk.KVStore, height exported.Height) {
+	key := host.ConsensusStateKey(height)
+	clientStore.Delete(key)
+}
+
 // IterateProcessedTime iterates through the prefix store and applies the callback.
 // If the cb returns true, then iterator will close and stop.
 func IterateProcessedTime(store sdk.KVStore, cb func(key, val []byte) bool) {
@@ -94,3 +124,123 @@ func GetProcessedTime(clientStore sdk.KVStore, height exported.Height) (uint64,
 	}
 	return sdk.BigEndianToUint64(bz), true
 }
+
+// deleteProcessedTime deletes the processedTime for a given height
+func deleteProcessedTime(clientStore sdk.KVStore, height exported.Height) {
+	key := ProcessedTimeKey(height)
+	clientStore.Delete(key)
+}
+
+// Iteration Code
+
+// IterationKey returns the key under which the consensus state key will be stored.
+// The iteration key is a BigEndian representation of the consensus state key to support efficient iteration.
+func IterationKey(height exported.Height) []byte {
+	heightBytes := bigEndianHeightBytes(height)
+	return append([]byte(KeyIterateConsensusStatePrefix), heightBytes...)
+}
+
+// SetIterationKey stores the consensus state key under a key that is more efficient for ordered iteration
+func SetIterationKey(clientStore sdk.KVStore, height exported.Height) {
+	key := IterationKey(height)
+	val := host.ConsensusStateKey(height)
+	clientStore.Set(key, val)
+}
+
+// GetIterationKey returns the consensus state key stored under the efficient iteration key.
+// NOTE: This function is currently only used for testing purposes
+func GetIterationKey(clientStore sdk.KVStore, height exported.Height) []byte {
+	key := IterationKey(height)
+	return clientStore.Get(key)
+}
+
+// deleteIterationKey deletes the iteration key for a given height
+func deleteIterationKey(clientStore sdk.KVStore, height exported.Height) {
+	key := IterationKey(height)
+	clientStore.Delete(key)
+}
+
+// GetHeightFromIterationKey takes an iteration key and returns the height that it references
+func GetHeightFromIterationKey(iterKey []byte) exported.Height {
+	bigEndianBytes := iterKey[len([]byte(KeyIterateConsensusStatePrefix)):]
+	revisionBytes := bigEndianBytes[0:8]
+	heightBytes := bigEndianBytes[8:]
+	revision := binary.BigEndian.Uint64(revisionBytes)
+	height := binary.BigEndian.Uint64(heightBytes)
+	return clienttypes.NewHeight(revision, height)
+}
+
+func IterateConsensusStateAscending(clientStore sdk.KVStore, cb func(height exported.Height) (stop bool)) error {
+	iterator := sdk.KVStorePrefixIterator(clientStore, []byte(KeyIterateConsensusStatePrefix))
+	defer iterator.Close()
+
+	for ; iterator.Valid(); iterator.Next() {
+		iterKey := iterator.Key()
+		height := GetHeightFromIterationKey(iterKey)
+		if cb(height) {
+			return nil
+		}
+	}
+	return nil
+}
+
+// GetNextConsensusState returns the lowest consensus state that is larger than the given height.
+// The Iterator returns a storetypes.Iterator which iterates from start (inclusive) to end (exclusive).
+// Thus, to get the next consensus state, we must first call iterator.Next() and then get the value.
+func GetNextConsensusState(clientStore sdk.KVStore, cdc codec.BinaryMarshaler, height exported.Height) (*ConsensusState, bool) {
+	iterateStore := prefix.NewStore(clientStore, []byte(KeyIterateConsensusStatePrefix))
+	iterator := iterateStore.Iterator(bigEndianHeightBytes(height), nil)
+	defer iterator.Close()
+	// ignore the consensus state at current height and get next height
+	iterator.Next()
+	if !iterator.Valid() {
+		return nil, false
+	}
+
+	csKey := iterator.Value()
+
+	return getTmConsensusState(clientStore, cdc, csKey)
+}
+
+// GetPreviousConsensusState returns the highest consensus state that is lower than the given height.
+// The Iterator returns a storetypes.Iterator which iterates from the end (exclusive) to start (inclusive).
+// Thus to get previous consensus state we call iterator.Value() immediately.
+func GetPreviousConsensusState(clientStore sdk.KVStore, cdc codec.BinaryMarshaler, height exported.Height) (*ConsensusState, bool) {
+	iterateStore := prefix.NewStore(clientStore, []byte(KeyIterateConsensusStatePrefix))
+	iterator := iterateStore.ReverseIterator(nil, bigEndianHeightBytes(height))
+	defer iterator.Close()
+
+	if !iterator.Valid() {
+		return nil, false
+	}
+
+	csKey := iterator.Value()
+
+	return getTmConsensusState(clientStore, cdc, csKey)
+}
+
+// Helper function for GetNextConsensusState and GetPreviousConsensusState
+func getTmConsensusState(clientStore sdk.KVStore, cdc codec.BinaryMarshaler, key []byte) (*ConsensusState, bool) {
+	bz := clientStore.Get(key)
+	if bz == nil {
+		return nil, false
+	}
+
+	consensusStateI, err := clienttypes.UnmarshalConsensusState(cdc, bz)
+	if err != nil {
+		return nil, false
+	}
+
+	consensusState, ok := consensusStateI.(*ConsensusState)
+	if !ok {
+		return nil, false
+	}
+	return consensusState, true
+}
+
+func bigEndianHeightBytes(height exported.Height) []byte {
+	heightBytes := make([]byte, 16)
+	binary.BigEndian.PutUint64(heightBytes, height.GetRevisionNumber())
+	binary.BigEndian.PutUint64(heightBytes[8:], height.GetRevisionHeight())
+	return heightBytes
+}

modules/light-clients/07-tendermint/types/store_test.go

@@ -1,7 +1,11 @@
 package types_test
 
 import (
+	"math"
+	"time"
+
 	clienttypes "github.com/cosmos/ibc-go/modules/core/02-client/types"
+	commitmenttypes "github.com/cosmos/ibc-go/modules/core/23-commitment/types"
 	host "github.com/cosmos/ibc-go/modules/core/24-host"
 	"github.com/cosmos/ibc-go/modules/core/exported"
 	solomachinetypes "github.com/cosmos/ibc-go/modules/light-clients/06-solomachine/types"
@@ -116,3 +120,75 @@ func (suite *TendermintTestSuite) TestGetProcessedTime() {
 	_, ok = types.GetProcessedTime(store, clienttypes.NewHeight(1, 1))
 	suite.Require().False(ok, "retrieved processed time for a non-existent consensus state")
 }
+
+func (suite *TendermintTestSuite) TestIterationKey() {
+	testHeights := []exported.Height{
+		clienttypes.NewHeight(0, 1),
+		clienttypes.NewHeight(0, 1234),
+		clienttypes.NewHeight(7890, 4321),
+		clienttypes.NewHeight(math.MaxUint64, math.MaxUint64),
+	}
+	for _, h := range testHeights {
+		k := types.IterationKey(h)
+		retrievedHeight := types.GetHeightFromIterationKey(k)
+		suite.Require().Equal(h, retrievedHeight, "retrieving height from iteration key failed")
+	}
+}
+
+func (suite *TendermintTestSuite) TestIterateConsensusStates() {
+	nextValsHash := []byte("nextVals")
+
+	// Set iteration keys and consensus states
+	types.SetIterationKey(suite.chainA.App.GetIBCKeeper().ClientKeeper.ClientStore(suite.chainA.GetContext(), "testClient"), clienttypes.NewHeight(0, 1))
+	suite.chainA.App.GetIBCKeeper().ClientKeeper.SetClientConsensusState(suite.chainA.GetContext(), "testClient", clienttypes.NewHeight(0, 1), types.NewConsensusState(time.Now(), commitmenttypes.NewMerkleRoot([]byte("hash0-1")), nextValsHash))
+	types.SetIterationKey(suite.chainA.App.GetIBCKeeper().ClientKeeper.ClientStore(suite.chainA.GetContext(), "testClient"), clienttypes.NewHeight(4, 9))
+	suite.chainA.App.GetIBCKeeper().ClientKeeper.SetClientConsensusState(suite.chainA.GetContext(), "testClient", clienttypes.NewHeight(4, 9), types.NewConsensusState(time.Now(), commitmenttypes.NewMerkleRoot([]byte("hash4-9")), nextValsHash))
+	types.SetIterationKey(suite.chainA.App.GetIBCKeeper().ClientKeeper.ClientStore(suite.chainA.GetContext(), "testClient"), clienttypes.NewHeight(0, 10))
+	suite.chainA.App.GetIBCKeeper().ClientKeeper.SetClientConsensusState(suite.chainA.GetContext(), "testClient", clienttypes.NewHeight(0, 10), types.NewConsensusState(time.Now(), commitmenttypes.NewMerkleRoot([]byte("hash0-10")), nextValsHash))
+	types.SetIterationKey(suite.chainA.App.GetIBCKeeper().ClientKeeper.ClientStore(suite.chainA.GetContext(), "testClient"), clienttypes.NewHeight(0, 4))
+	suite.chainA.App.GetIBCKeeper().ClientKeeper.SetClientConsensusState(suite.chainA.GetContext(), "testClient", clienttypes.NewHeight(0, 4), types.NewConsensusState(time.Now(), commitmenttypes.NewMerkleRoot([]byte("hash0-4")), nextValsHash))
+	types.SetIterationKey(suite.chainA.App.GetIBCKeeper().ClientKeeper.ClientStore(suite.chainA.GetContext(), "testClient"), clienttypes.NewHeight(40, 1))
+	suite.chainA.App.GetIBCKeeper().ClientKeeper.SetClientConsensusState(suite.chainA.GetContext(), "testClient", clienttypes.NewHeight(40, 1), types.NewConsensusState(time.Now(), commitmenttypes.NewMerkleRoot([]byte("hash40-1")), nextValsHash))
+
+	var testArr []string
+	cb := func(height exported.Height) bool {
+		testArr = append(testArr, height.String())
+		return false
+	}
+
+	types.IterateConsensusStateAscending(suite.chainA.App.GetIBCKeeper().ClientKeeper.ClientStore(suite.chainA.GetContext(), "testClient"), cb)
+	expectedArr := []string{"0-1", "0-4", "0-10", "4-9", "40-1"}
+	suite.Require().Equal(expectedArr, testArr)
+}
+
+func (suite *TendermintTestSuite) TestGetNeighboringConsensusStates() {
+	nextValsHash := []byte("nextVals")
+	cs01 := types.NewConsensusState(time.Now().UTC(), commitmenttypes.NewMerkleRoot([]byte("hash0-1")), nextValsHash)
+	cs04 := types.NewConsensusState(time.Now().UTC(), commitmenttypes.NewMerkleRoot([]byte("hash0-4")), nextValsHash)
+	cs49 := types.NewConsensusState(time.Now().UTC(), commitmenttypes.NewMerkleRoot([]byte("hash4-9")), nextValsHash)
+	height01 := clienttypes.NewHeight(0, 1)
+	height04 := clienttypes.NewHeight(0, 4)
+	height49 := clienttypes.NewHeight(4, 9)
+
+	// Set iteration keys and consensus states
+	types.SetIterationKey(suite.chainA.App.GetIBCKeeper().ClientKeeper.ClientStore(suite.chainA.GetContext(), "testClient"), height01)
+	suite.chainA.App.GetIBCKeeper().ClientKeeper.SetClientConsensusState(suite.chainA.GetContext(), "testClient", height01, cs01)
+	types.SetIterationKey(suite.chainA.App.GetIBCKeeper().ClientKeeper.ClientStore(suite.chainA.GetContext(), "testClient"), height04)
+	suite.chainA.App.GetIBCKeeper().ClientKeeper.SetClientConsensusState(suite.chainA.GetContext(), "testClient", height04, cs04)
+	types.SetIterationKey(suite.chainA.App.GetIBCKeeper().ClientKeeper.ClientStore(suite.chainA.GetContext(), "testClient"), height49)
+	suite.chainA.App.GetIBCKeeper().ClientKeeper.SetClientConsensusState(suite.chainA.GetContext(), "testClient", height49, cs49)
+
+	prevCs01, ok := types.GetPreviousConsensusState(suite.chainA.App.GetIBCKeeper().ClientKeeper.ClientStore(suite.chainA.GetContext(), "testClient"), suite.chainA.Codec, height01)
+	suite.Require().Nil(prevCs01, "consensus state exists before lowest consensus state")
+	suite.Require().False(ok)
+	prevCs49, ok := types.GetPreviousConsensusState(suite.chainA.App.GetIBCKeeper().ClientKeeper.ClientStore(suite.chainA.GetContext(), "testClient"), suite.chainA.Codec, height49)
+	suite.Require().Equal(cs04, prevCs49, "previous consensus state is not returned correctly")
+	suite.Require().True(ok)
+
+	nextCs01, ok := types.GetNextConsensusState(suite.chainA.App.GetIBCKeeper().ClientKeeper.ClientStore(suite.chainA.GetContext(), "testClient"), suite.chainA.Codec, height01)
+	suite.Require().Equal(cs04, nextCs01, "next consensus state not returned correctly")
+	suite.Require().True(ok)
+	nextCs49, ok := types.GetNextConsensusState(suite.chainA.App.GetIBCKeeper().ClientKeeper.ClientStore(suite.chainA.GetContext(), "testClient"), suite.chainA.Codec, height49)
+	suite.Require().Nil(nextCs49, "next consensus state exists after highest consensus state")
+	suite.Require().False(ok)
+}

modules/light-clients/07-tendermint/types/update.go

@@ -37,6 +37,11 @@ import (
 // number must be the same. To update to a new revision, use a separate upgrade path
 // Tendermint client validity checking uses the bisection algorithm described
 // in the [Tendermint spec](https://github.com/tendermint/spec/blob/master/spec/consensus/light-client.md).
+//
+// Pruning:
+// UpdateClient will additionally retrieve the earliest consensus state for this clientID and check if it is expired. If it is,
+// that consensus state will be pruned from store along with all associated metadata. This will prevent the client store from
+// becoming bloated with expired consensus states that can no longer be used for updates and packet verification.
 func (cs ClientState) CheckHeaderAndUpdateState(
 	ctx sdk.Context, cdc codec.BinaryMarshaler, clientStore sdk.KVStore,
 	header exported.Header,
@@ -60,6 +65,35 @@ func (cs ClientState) CheckHeaderAndUpdateState(
 		return nil, nil, err
 	}
 
+	// Check the earliest consensus state to see if it is expired, if so then set the prune height
+	// so that we can delete consensus state and all associated metadata.
+	var (
+		pruneHeight exported.Height
+		pruneError  error
+	)
+	pruneCb := func(height exported.Height) bool {
+		consState, err := GetConsensusState(clientStore, cdc, height)
+		// this error should never occur
+		if err != nil {
+			pruneError = err
+			return true
+		}
+		if cs.IsExpired(consState.Timestamp, ctx.BlockTime()) {
+			pruneHeight = height
+		}
+		return true
+	}
+	IterateConsensusStateAscending(clientStore, pruneCb)
+	if pruneError != nil {
+		return nil, nil, pruneError
+	}
+	// if pruneHeight is set, delete consensus state and metadata
+	if pruneHeight != nil {
+		deleteConsensusState(clientStore, pruneHeight)
+		deleteProcessedTime(clientStore, pruneHeight)
+		deleteIterationKey(clientStore, pruneHeight)
+	}
+
 	newClientState, consensusState := update(ctx, clientStore, &cs, tmHeader)
 	return newClientState, consensusState, nil
 }
@@ -180,7 +214,9 @@ func update(ctx sdk.Context, clientStore sdk.KVStore, clientState *ClientState,
 
 	// set context time as processed time as this is state internal to tendermint client logic.
 	// client state and consensus state will be set by client keeper
+	// set iteration key to provide ability for efficient ordered iteration of consensus states.
 	SetProcessedTime(clientStore, header.GetHeight(), uint64(ctx.BlockTime().UnixNano()))
+	SetIterationKey(clientStore, header.GetHeight())
 
 	return clientState, consensusState
 }