use an address <-> zone mapping to avoid heavy lookups (#394)

x/interchainstaking/keeper/ibc_channel_handlers.go

@@ -51,6 +51,8 @@ func (k *Keeper) HandleChannelOpenAck(ctx sdk.Context, portID, connectionID stri
 				return err
 			}
 
+			k.SetAddressZoneMapping(ctx, address, zone.ChainId)
+
 			balanceQuery := bankTypes.QueryAllBalancesRequest{Address: address}
 			bz, err := k.GetCodec().Marshal(&balanceQuery)
 			if err != nil {
@@ -77,6 +79,7 @@ func (k *Keeper) HandleChannelOpenAck(ctx sdk.Context, portID, connectionID stri
 			if err != nil {
 				return err
 			}
+			k.SetAddressZoneMapping(ctx, address, zone.ChainId)
 		}
 
 	// delegation addresses
@@ -86,9 +89,13 @@ func (k *Keeper) HandleChannelOpenAck(ctx sdk.Context, portID, connectionID stri
 			if err != nil {
 				return err
 			}
+
+			k.SetAddressZoneMapping(ctx, address, zone.ChainId)
+
 			if _, err = k.FlushOutstandingDelegations(ctx, &zone); err != nil {
 				k.Logger(ctx).Error("unable to clear outstanding delegations", "error", err)
 			}
+
 		}
 
 	// performance address
@@ -99,6 +106,7 @@ func (k *Keeper) HandleChannelOpenAck(ctx sdk.Context, portID, connectionID stri
 			if err != nil {
 				return err
 			}
+			k.SetAddressZoneMapping(ctx, address, zone.ChainId)
 		}
 
 		// emit this periodic query the first time, but not subsequently.

x/interchainstaking/keeper/zones.go

@@ -4,7 +4,6 @@ import (
 	"errors"
 	"fmt"
 	"math"
-	"strings"
 
 	"github.com/cosmos/cosmos-sdk/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -13,8 +12,6 @@ import (
 	distrTypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 
-	icatypes "github.com/cosmos/ibc-go/v5/modules/apps/27-interchain-accounts/types"
-
 	icqtypes "github.com/ingenuity-build/quicksilver/x/interchainquery/types"
 	"github.com/ingenuity-build/quicksilver/x/interchainstaking/types"
 )
@@ -67,6 +64,28 @@ func (k *Keeper) IterateZones(ctx sdk.Context, fn func(index int64, zoneInfo *ty
 	}
 }
 
+// GetAddressZoneMapping returns zone <-> address mapping.
+func (k *Keeper) GetAddressZoneMapping(ctx sdk.Context, address string) (string, bool) {
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixAddressZoneMapping)
+	bz := store.Get([]byte(address))
+	if len(bz) == 0 {
+		return "", false
+	}
+	return string(bz), true
+}
+
+// SetAddressZoneMapping set zone <-> address mapping.
+func (k *Keeper) SetAddressZoneMapping(ctx sdk.Context, address, chainID string) {
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixAddressZoneMapping)
+	store.Set([]byte(address), []byte(chainID))
+}
+
+// DeleteAddressZoneMapping delete zone info.
+func (k *Keeper) DeleteAddressZoneMapping(ctx sdk.Context, address string) {
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixAddressZoneMapping)
+	store.Delete([]byte(address))
+}
+
 func (k *Keeper) GetDelegatedAmount(ctx sdk.Context, zone *types.Zone) sdk.Coin {
 	out := sdk.NewCoin(zone.BaseDenom, sdk.ZeroInt())
 	k.IterateAllDelegations(ctx, zone, func(delegation types.Delegation) (stop bool) {
@@ -110,73 +129,39 @@ func (k *Keeper) GetZoneFromContext(ctx sdk.Context) (*types.Zone, error) {
 	return &zone, nil
 }
 
+func (k *Keeper) GetZoneForAccount(ctx sdk.Context, address string) (zone *types.Zone) {
+	chainID, found := k.GetAddressZoneMapping(ctx, address)
+	if !found {
+		return nil
+	}
+	// doesn't matter if this is _found_ because we are expecting to return nil if not anyway.
+	z, _ := k.GetZone(ctx, chainID)
+	return &z
+}
+
 // GetZoneForDelegateAccount determines the zone for a given address.
 func (k *Keeper) GetZoneForDelegateAccount(ctx sdk.Context, address string) *types.Zone {
-	var zone *types.Zone
-	k.IterateZones(ctx, func(_ int64, zoneInfo *types.Zone) (stop bool) {
-		if zoneInfo.DelegationAddress != nil && zoneInfo.DelegationAddress.Address == address {
-			zone = zoneInfo
-			return true
-		}
-		return false
-	})
-	return zone
+	z := k.GetZoneForAccount(ctx, address)
+	if z != nil && z.DelegationAddress != nil && address == z.DelegationAddress.Address {
+		return z
+	}
+	return nil
 }
 
 func (k *Keeper) GetZoneForPerformanceAccount(ctx sdk.Context, address string) *types.Zone {
-	var zone *types.Zone
-	k.IterateZones(ctx, func(_ int64, zoneInfo *types.Zone) (stop bool) {
-		if zoneInfo.PerformanceAddress != nil && zoneInfo.PerformanceAddress.Address == address {
-			zone = zoneInfo
-			return true
-		}
-		return false
-	})
-	return zone
-}
-
-func (k *Keeper) GetZoneForDepositAccount(ctx sdk.Context, address string) *types.Zone {
-	var zone *types.Zone
-	k.IterateZones(ctx, func(_ int64, zoneInfo *types.Zone) (stop bool) {
-		if zoneInfo.DepositAddress != nil && zoneInfo.DepositAddress.Address == address {
-			zone = zoneInfo
-			return true
-		}
-		return false
-	})
-	return zone
-}
-
-func (k *Keeper) EnsureICAsActive(ctx sdk.Context, zone *types.Zone) error {
-	k.Logger(ctx).Info("Ensuring ICAs for zone", "zone", zone.ChainId)
-	if err := k.EnsureICAActive(ctx, zone, zone.DepositAddress); err != nil {
-		return err
+	z := k.GetZoneForAccount(ctx, address)
+	if z != nil && z.PerformanceAddress != nil && address == z.PerformanceAddress.Address {
+		return z
 	}
-	if err := k.EnsureICAActive(ctx, zone, zone.DelegationAddress); err != nil {
-		return err
-	}
-	if err := k.EnsureICAActive(ctx, zone, zone.PerformanceAddress); err != nil {
-		return err
-	}
-	return k.EnsureICAActive(ctx, zone, zone.WithdrawalAddress)
+	return nil
 }
 
-func (k *Keeper) EnsureICAActive(ctx sdk.Context, zone *types.Zone, account *types.ICAAccount) error {
-	if account == nil {
-		k.Logger(ctx).Info("Account does not exist")
-		// address has not been set yet. nothing to check.
-		return nil
-	}
-
-	if _, found := k.ICAControllerKeeper.GetOpenActiveChannel(ctx, zone.ConnectionId, account.GetPortName()); found {
-		k.Logger(ctx).Info("Account is active", "account", account.Address)
-		// channel is active. all is well :)
-		return nil
+func (k *Keeper) GetZoneForDepositAccount(ctx sdk.Context, address string) *types.Zone {
+	z := k.GetZoneForAccount(ctx, address)
+	if z != nil && z.DepositAddress != nil && address == z.DepositAddress.Address {
+		return z
 	}
-
-	// channel is not active; attempt reopen.
-	k.Logger(ctx).Error("channel is inactive. attempting to reopen.", "connection", zone.ConnectionId, "port", account.GetPortName())
-	return k.ICAControllerKeeper.RegisterInterchainAccount(ctx, zone.ConnectionId, strings.TrimPrefix(account.GetPortName(), icatypes.PortPrefix), "")
+	return nil
 }
 
 func (k *Keeper) EnsureWithdrawalAddresses(ctx sdk.Context, zone *types.Zone) error {

x/interchainstaking/keeper/zones_test.go

@@ -15,6 +15,7 @@ import (
 	"golang.org/x/exp/maps"
 
 	"github.com/ingenuity-build/quicksilver/app"
+	"github.com/ingenuity-build/quicksilver/utils"
 	"github.com/ingenuity-build/quicksilver/x/interchainstaking/types"
 )
 
@@ -73,20 +74,22 @@ func TestKeeperWithZonesRoundTrip(t *testing.T) {
 	chainIDPrefix := "quicksilver-"
 	indexToZone := make(map[int64]types.Zone, nzones)
 	for i := 0; i < nzones; i++ {
-		chainID := fmt.Sprintf("%s-%d", chainIDPrefix, i)
+		chainID := fmt.Sprintf("%s%d", chainIDPrefix, i)
+		delegationAddr := utils.GenerateAccAddressForTestWithPrefix("cosmos")
 		zone := types.Zone{
 			ConnectionId: "conn-test",
 			ChainId:      chainID,
 			LocalDenom:   "qck",
 			BaseDenom:    "qck",
 			DelegationAddress: &types.ICAAccount{
-				Address: "cosmos1ssrxxe4xsls57ehrkswlkhlkcverf0p0fpgyhzqw0hfdqj92ynxsw29r6e",
+				Address: delegationAddr,
 				Balance: sdk.NewCoins(
 					sdk.NewCoin("qck", sdk.NewInt(100)),
 					sdk.NewCoin("uqck", sdk.NewInt(700000)),
 				),
 			},
 		}
+		kpr.SetAddressZoneMapping(ctx, delegationAddr, zone.ChainId)
 		kpr.SetZone(ctx, &zone)
 		gotZone, ok := kpr.GetZone(ctx, chainID)
 		require.True(t, ok, "expected to retrieve the correct zone")
@@ -132,6 +135,7 @@ func TestKeeperWithZonesRoundTrip(t *testing.T) {
 			sdk.NewCoin("uqck", sdk.NewInt(900000)),
 		),
 	}
+	kpr.SetAddressZoneMapping(ctx, "cosmosvaloper1sjllsnramtg3ewxqwwrwjxfgc4n4ef9u2lcnj0", perfAcctZone.ChainId)
 	kpr.SetZone(ctx, &perfAcctZone)
 	gotPerfAcctZone := kpr.GetZoneForPerformanceAccount(ctx, perfAcctZone.PerformanceAddress.Address)
 	require.Equal(t, &perfAcctZone, gotPerfAcctZone, "expecting a match in performance accounts")
@@ -152,10 +156,11 @@ func TestKeeperWithZonesRoundTrip(t *testing.T) {
 	// 7.2. Set some delegations.
 	del1 := types.Delegation{
 		Amount:            sdk.NewCoin(firstZone.BaseDenom, sdk.NewInt(17000)),
-		DelegationAddress: "cosmos1ssrxxe4xsls57ehrkswlkhlkcverf0p0fpgyhzqw0hfdqj92ynxsw29r6e",
+		DelegationAddress: firstZone.DelegationAddress.Address,
 		Height:            10,
 		ValidatorAddress:  "cosmosvaloper1sjllsnramtg3ewxqwwrwjxfgc4n4ef9u2lcnj0",
 	}
+
 	kpr.SetDelegation(ctx, &firstZone, del1)
 
 	// 7.3. Retrieve the delegation now, it should be set.

x/interchainstaking/types/keys.go

@@ -54,7 +54,9 @@ var (
 	KeyPrefixPerformanceDelegation       = []byte{0x08}
 	KeyPrefixSnapshotIntent              = []byte{0x09}
 	KeyPrefixRequeuedWithdrawalRecordSeq = []byte{0x0a}
-	// fill in missing 0b - 0f before adding 0x11!
+	KeyPrefixAddressZoneMapping          = []byte{0x0b}
+
+	// fill in missing 0c - 0f before adding 0x11!
 	KeyPrefixRedelegationRecord = []byte{0x10}
 )