fix: fix bugs in calculation of pool-service total delegator shares (#38

)

## Description

This PR fixes a bug that pool-service total delegator shares were
miscalculated by ignoring user preferences when a user modifies the
delegation.

This PR also implements the logic that clears total shares state when a
service is deleted.

Also, with #26 we changed the user preferences but didn't update
`PoolServiceTotalDelegatorShares` accordingly. Normally this should've
been done by hooks but since we don't call hooks in store migrations we
need to do this manually. The best place to do this would be inside the
fork handler.

Please note that we might need to apply the same thing when upgrading
the testnet in the future.

<!-- Add a description of the changes that this PR introduces and the
files that
are the most critical to review. -->

---

### Author Checklist

*All items are required. Please add a note to the item if the item is
not applicable and
please add links to any relevant follow up issues.*

I have...

- [ ] included the correct [type
prefix](https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.json)
in the PR title
- [ ] added `!` to the type prefix if API or client breaking change
- [ ] targeted the correct branch (see [PR
Targeting](https://github.com/milkyway-labs/milkyway/blob/master/CONTRIBUTING.md#pr-targeting))
- [ ] provided a link to the relevant issue or specification
- [ ] followed the guidelines for [building
modules](https://docs.cosmos.network/v0.44/building-modules/intro.html)
- [ ] included the necessary unit and integration
[tests](https://github.com/milkyway-labs/milkyway/blob/master/CONTRIBUTING.md#testing)
- [ ] added a changelog entry to `CHANGELOG.md`
- [ ] included comments for [documenting Go
code](https://blog.golang.org/godoc)
- [ ] updated the relevant documentation or specification
- [ ] reviewed "Files changed" and left comments if necessary
- [ ] confirmed all CI checks have passed

### Reviewers Checklist

*All items are required. Please add a note if the item is not applicable
and please add
your handle next to the items reviewed if you only reviewed selected
items.*

I have...

- [ ] confirmed the correct [type
prefix](https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.json)
in the PR title
- [ ] confirmed `!` in the type prefix if API or client breaking change
- [ ] confirmed all author checklist items have been addressed
- [ ] reviewed state machine logic
- [ ] reviewed API design and naming
- [ ] reviewed documentation is accurate
- [ ] reviewed tests and test coverage
- [ ] manually tested (if applicable)

(cherry picked from commit 1addfbe)

app/forks/v8/fork.go

@@ -7,6 +7,7 @@ import (
 	"github.com/milkyway-labs/milkyway/v7/app/keepers"
 	"github.com/milkyway-labs/milkyway/v7/utils"
 	poolstypes "github.com/milkyway-labs/milkyway/v7/x/pools/types"
+	restakingtypes "github.com/milkyway-labs/milkyway/v7/x/restaking/types"
 	servicestypes "github.com/milkyway-labs/milkyway/v7/x/services/types"
 )
 
@@ -31,19 +32,31 @@ func BeginFork(ctx sdk.Context, mm *module.Manager, cfg module.Configurator, kee
 		panic(err)
 	}
 
+	err = updateServiceSecuringPools(ctx, keepers)
+	if err != nil {
+		panic(err)
+	}
+
+	err = updatePoolServiceTotalDelShares(ctx, keepers)
+	if err != nil {
+		panic(err)
+	}
+}
+
+func updateServiceSecuringPools(ctx sdk.Context, keepers *keepers.AppKeepers) error {
 	// Get all pools IDs first.
 	pools, err := keepers.PoolsKeeper.GetPools(ctx)
 	if err != nil {
-		panic(err)
+		return err
 	}
-	poolIDs := utils.Map(pools, func(pool poolstypes.Pool) uint32 {
+	poolsIDs := utils.Map(pools, func(pool poolstypes.Pool) uint32 {
 		return pool.ID
 	})
 
 	// Update existing services that have no securing pools configured(which means
 	// they are secured by all pools by default) to have all the pools as their
 	// securing pools.
-	var serviceIDsToUpdate []uint32
+	var servicesIDsToUpdate []uint32
 	err = keepers.ServicesKeeper.IterateServices(ctx, func(service servicestypes.Service) (stop bool, err error) {
 		// We only need to update services that have not been configured to secure pools
 		// since we changed the semantic of "empty securing pools"
@@ -55,19 +68,75 @@ func BeginFork(ctx sdk.Context, mm *module.Manager, cfg module.Configurator, kee
 			return false, nil
 		}
 
-		serviceIDsToUpdate = append(serviceIDsToUpdate, service.ID)
+		servicesIDsToUpdate = append(servicesIDsToUpdate, service.ID)
 		return false, nil
 	})
 	if err != nil {
-		panic(err)
+		return err
 	}
 
-	for _, serviceID := range serviceIDsToUpdate {
-		for _, poolID := range poolIDs {
+	for _, serviceID := range servicesIDsToUpdate {
+		for _, poolID := range poolsIDs {
 			err = keepers.RestakingKeeper.AddPoolToServiceSecuringPools(ctx, serviceID, poolID)
 			if err != nil {
-				panic(err)
+				return err
 			}
 		}
 	}
+
+	return nil
+}
+
+// updatePoolServiceTotalDelShares synchronizes the pool-service total delegator shares
+// with the updated user preferences.
+func updatePoolServiceTotalDelShares(ctx sdk.Context, keepers *keepers.AppKeepers) error {
+	preferencesCache := map[string]restakingtypes.UserPreferences{}
+
+	allServices, err := keepers.ServicesKeeper.GetServices(ctx)
+	if err != nil {
+		return err
+	}
+	allServicesIDs := utils.Map(allServices, func(service servicestypes.Service) uint32 {
+		return service.ID
+	})
+
+	// First clear all the existing pool-service total delegator shares.
+	err = keepers.RewardsKeeper.PoolServiceTotalDelegatorShares.Clear(ctx, nil)
+	if err != nil {
+		return err
+	}
+
+	// Iterating over all pool delegations, increment pool-service total delegator shares
+	// only if the service is trusted with the pool.
+	err = keepers.RestakingKeeper.IterateAllPoolDelegations(ctx, func(delegation restakingtypes.Delegation) (stop bool, err error) {
+		preferences, ok := preferencesCache[delegation.UserAddress]
+		if !ok {
+			preferences, err = keepers.RestakingKeeper.GetUserPreferences(ctx, delegation.UserAddress)
+			if err != nil {
+				return true, err
+			}
+			preferencesCache[delegation.UserAddress] = preferences
+		}
+
+		for _, serviceID := range allServicesIDs {
+			if !preferences.IsServiceTrustedWithPool(serviceID, delegation.TargetID) {
+				continue
+			}
+
+			err = keepers.RewardsKeeper.IncrementPoolServiceTotalDelegatorShares(
+				ctx,
+				delegation.TargetID,
+				serviceID,
+				delegation.Shares,
+			)
+			if err != nil {
+				return true, err
+			}
+		}
+		return false, nil
+	})
+	if err != nil {
+		return err
+	}
+	return nil
 }

x/rewards/keeper/allocation_test.go

@@ -474,7 +474,7 @@ func (suite *KeeperTestSuite) TestAllocateRewards_TransferRewardsOnlyWhenAllocat
 	balances = suite.bankKeeper.GetAllBalances(ctx, rewardsPool)
 	suite.Require().Empty(balances)
 
-	// Finally, Alice trust the service through the pool and rewards are allocated.
+	// Finally, Alice trusts the service through the pool and rewards are allocated.
 	suite.SetUserPreferences(ctx, aliceAddr.String(), []restakingtypes.TrustedServiceEntry{
 		restakingtypes.NewTrustedServiceEntry(service.ID, nil),
 	})
@@ -598,6 +598,61 @@ func (suite *KeeperTestSuite) TestAllocateRewards_TransferRewardsOnlyWhenAllocat
 	suite.Require().NotEmpty(balances)
 }
 
+func (suite *KeeperTestSuite) TestAllocateRewards_NoRewardsForNonTrustedPools() {
+	// Cache the context to avoid test conflicts
+	ctx, _ := suite.ctx.CacheContext()
+
+	suite.RegisterCurrency(ctx, "umilk", "MILK", 6, utils.MustParseDec("1"))
+
+	// Create a service.
+	serviceAdmin := testutil.TestAddress(10000)
+	service := suite.CreateService(ctx, "Service", serviceAdmin.String())
+
+	// Create an active rewards plan.
+	suite.CreateBasicRewardsPlan(
+		ctx,
+		service.ID,
+		utils.MustParseCoin("1000_000000service"),
+		time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC),
+		time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC),
+		utils.MustParseCoins("100000_000000service"),
+	)
+	suite.AddPoolsToServiceSecuringPools(ctx, service.ID, []uint32{1})
+
+	// Call AllocateRewards to set last rewards allocation time.
+	err := suite.keeper.AllocateRewards(ctx)
+	suite.Require().NoError(err)
+
+	// Alice, Bob and Charlie delegate $MILK to the pool.
+	aliceAddr := testutil.TestAddress(1)
+	suite.DelegatePool(ctx, utils.MustParseCoin("10_000000umilk"), aliceAddr.String(), true)
+	bobAddr := testutil.TestAddress(2)
+	suite.DelegatePool(ctx, utils.MustParseCoin("10_000000umilk"), bobAddr.String(), true)
+	charlieAddr := testutil.TestAddress(3)
+	suite.DelegatePool(ctx, utils.MustParseCoin("10_000000umilk"), charlieAddr.String(), true)
+
+	// Alice trusts the service through pool 2, which doesn't exist.
+	suite.SetUserPreferences(ctx, aliceAddr.String(), []restakingtypes.TrustedServiceEntry{
+		restakingtypes.NewTrustedServiceEntry(1, []uint32{2}),
+	})
+	// Bob doesn't trust the service at all.
+	suite.SetUserPreferences(ctx, bobAddr.String(), []restakingtypes.TrustedServiceEntry{
+		restakingtypes.NewTrustedServiceEntry(2, nil),
+	})
+
+	// Only Charlie receives rewards.
+	ctx = suite.allocateRewards(ctx, 10*time.Second)
+	rewards, err := suite.keeper.GetPoolDelegationRewards(ctx, aliceAddr, 1)
+	suite.Require().NoError(err)
+	suite.Require().Empty(rewards)
+	rewards, err = suite.keeper.GetPoolDelegationRewards(ctx, bobAddr, 1)
+	suite.Require().NoError(err)
+	suite.Require().Empty(rewards)
+	rewards, err = suite.keeper.GetPoolDelegationRewards(ctx, charlieAddr, 1)
+	suite.Require().NoError(err)
+	suite.Require().Equal("115740.000000000000000000service", rewards.Sum().String())
+}
+
 func (suite *KeeperTestSuite) TestAllocateRewards_WeightedDistributions() {
 	// Cache the context to avoid test conflicts
 	ctx, _ := suite.ctx.CacheContext()

x/rewards/keeper/delegation.go

@@ -92,8 +92,14 @@ func (k *Keeper) calculateDelegationRewardsBetween(
 			return nil, err
 		}
 
+		preferences, err := k.restakingKeeper.GetUserPreferences(ctx, delegator)
+		if err != nil {
+			return nil, err
+		}
+
 		for _, diff := range differences {
-			if slices.Contains(servicesIDs, diff.ServiceID) {
+			if slices.Contains(servicesIDs, diff.ServiceID) &&
+				preferences.IsServiceTrustedWithPool(diff.ServiceID, target.GetID()) {
 				decPools = decPools.Add(diff.DecPools...)
 			}
 		}

x/rewards/keeper/hooks.go

@@ -3,10 +3,12 @@ package keeper
 import (
 	"context"
 
+	"cosmossdk.io/collections"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 
 	"github.com/milkyway-labs/milkyway/v7/utils"
 	restakingtypes "github.com/milkyway-labs/milkyway/v7/x/restaking/types"
+	"github.com/milkyway-labs/milkyway/v7/x/rewards/types"
 	servicestypes "github.com/milkyway-labs/milkyway/v7/x/services/types"
 )
 
@@ -65,19 +67,28 @@ func (k *Keeper) BeforeDelegationSharesModified(ctx context.Context, delType res
 	}
 
 	if delType == restakingtypes.DELEGATION_TYPE_POOL {
+		// Note that it'll return all services IDs that are stored when the user sets no
+		// entries in its preferences.
 		servicesIDs, err := k.restakingKeeper.GetUserTrustedServicesIDs(ctx, delegator)
 		if err != nil {
 			return err
 		}
 
+		preferences, err := k.restakingKeeper.GetUserPreferences(ctx, delegator)
+		if err != nil {
+			return err
+		}
+
 		for _, serviceID := range servicesIDs {
-			// We decrement the amount of shares within the pool-service pair here so that we
-			// can later increment those shares again within the AfterDelegationModified
-			// hook. This is due in order to keep consistency if the shares change due to a
-			// new delegation or an undelegation
-			err = k.DecrementPoolServiceTotalDelegatorShares(ctx, targetID, serviceID, del.Shares)
-			if err != nil {
-				return err
+			if preferences.IsServiceTrustedWithPool(serviceID, targetID) {
+				// We decrement the amount of shares within the pool-service pair here so that we
+				// can later increment those shares again within the AfterDelegationModified
+				// hook. This is due in order to keep consistency if the shares change due to a
+				// new delegation or an undelegation
+				err = k.DecrementPoolServiceTotalDelegatorShares(ctx, targetID, serviceID, del.Shares)
+				if err != nil {
+					return err
+				}
 			}
 		}
 	}
@@ -112,26 +123,65 @@ func (k *Keeper) AfterDelegationModified(ctx context.Context, delType restakingt
 			return sdkerrors.ErrNotFound.Wrapf("pool delegation not found: %d, %s", targetID, delegator)
 		}
 
+		// Note that it'll return all services IDs that are stored when the user sets no
+		// entries in its preferences.
 		servicesIDs, err := k.restakingKeeper.GetUserTrustedServicesIDs(ctx, delegator)
 		if err != nil {
 			return err
 		}
 
+		preferences, err := k.restakingKeeper.GetUserPreferences(ctx, delegator)
+		if err != nil {
+			return err
+		}
+
 		for _, serviceID := range servicesIDs {
-			// We decremented the amount of shares within the pool-service pair in the
-			// BeforeDelegationSharesModified hook. We increment the shares here again
-			// to keep consistency if the shares change due to a new delegation or an
-			// undelegation
-			err = k.IncrementPoolServiceTotalDelegatorShares(ctx, targetID, serviceID, delegation.Shares)
-			if err != nil {
-				return err
+			if preferences.IsServiceTrustedWithPool(serviceID, targetID) {
+				// We decremented the amount of shares within the pool-service pair in the
+				// BeforeDelegationSharesModified hook. We increment the shares here again
+				// to keep consistency if the shares change due to a new delegation or an
+				// undelegation
+				err = k.IncrementPoolServiceTotalDelegatorShares(ctx, targetID, serviceID, delegation.Shares)
+				if err != nil {
+					return err
+				}
 			}
 		}
 	}
 
 	return nil
 }
 
+// BeforeServiceDeleted is called before a service is deleted
+func (k *Keeper) BeforeServiceDeleted(ctx context.Context, serviceID uint32) error {
+	err := k.BeforeDelegationTargetRemoved(ctx, restakingtypes.DELEGATION_TYPE_SERVICE, serviceID)
+	if err != nil {
+		return err
+	}
+
+	// Clean up pool-service total delegator shares for the service
+	var keysToDelete []collections.Pair[uint32, uint32]
+	err = k.PoolServiceTotalDelegatorShares.Walk(ctx, nil, func(key collections.Pair[uint32, uint32], _ types.PoolServiceTotalDelegatorShares) (stop bool, err error) {
+		if key.K2() != serviceID {
+			return false, nil
+		}
+		keysToDelete = append(keysToDelete, key)
+		return false, nil
+	})
+	if err != nil {
+		return err
+	}
+
+	for _, key := range keysToDelete {
+		err = k.PoolServiceTotalDelegatorShares.Remove(ctx, key)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // AfterUserPreferencesModified is called after a user's trust in a service is
 // updated. It updates the total delegator shares for the service in all pools
 // where the user has a delegation.

x/rewards/keeper/hooks_test.go

@@ -0,0 +1,49 @@
+package keeper_test
+
+import (
+	"github.com/milkyway-labs/milkyway/v7/app/testutil"
+	"github.com/milkyway-labs/milkyway/v7/utils"
+	restakingtypes "github.com/milkyway-labs/milkyway/v7/x/restaking/types"
+)
+
+func (suite *KeeperTestSuite) TestPoolServiceTotalDelegatorShares() {
+	// Cache the context to avoid test conflicts
+	ctx, _ := suite.ctx.CacheContext()
+
+	suite.RegisterCurrency(ctx, "umilk", "umilk", 6, utils.MustParseDec("1"))
+
+	// Create a service.
+	serviceAdmin := testutil.TestAddress(10000)
+	service := suite.CreateService(ctx, "Service", serviceAdmin.String())
+
+	// Alice by default trusts all services through all pools, so when Alice delegates
+	// it increments pool-service total delegator shares of the pool.
+	aliceAddr := testutil.TestAddress(1)
+	suite.DelegatePool(ctx, utils.MustParseCoin("10_000000umilk"), aliceAddr.String(), true)
+
+	shares, err := suite.keeper.GetPoolServiceTotalDelegatorShares(ctx, 1, service.ID)
+	suite.Require().NoError(err)
+	suite.Require().Equal(utils.MustParseDecCoins("10_000000pool/1/umilk"), shares)
+
+	// Bob doesn't trust the service through the pool, so when Bob delegates it
+	// doesn't increment pool-service total delegator shares of the pool.
+	bobAddr := testutil.TestAddress(2)
+	suite.SetUserPreferences(ctx, bobAddr.String(), []restakingtypes.TrustedServiceEntry{
+		restakingtypes.NewTrustedServiceEntry(1, []uint32{2}),
+	})
+	suite.DelegatePool(ctx, utils.MustParseCoin("10_000000umilk"), bobAddr.String(), true)
+
+	shares, err = suite.keeper.GetPoolServiceTotalDelegatorShares(ctx, 1, service.ID)
+	suite.Require().NoError(err)
+	suite.Require().Equal(utils.MustParseDecCoins("10_000000pool/1/umilk"), shares)
+
+	// But when Bob decides to trust the service through the pool, it increments
+	// the total shares.
+
+	suite.SetUserPreferences(ctx, bobAddr.String(), []restakingtypes.TrustedServiceEntry{
+		restakingtypes.NewTrustedServiceEntry(1, []uint32{service.ID}),
+	})
+	shares, err = suite.keeper.GetPoolServiceTotalDelegatorShares(ctx, 1, service.ID)
+	suite.Require().NoError(err)
+	suite.Require().Equal(utils.MustParseDecCoins("20_000000pool/1/umilk"), shares)
+}

x/rewards/keeper/services_hooks.go

@@ -34,7 +34,7 @@ func (h ServicesHooks) AfterServiceDeactivated(context.Context, uint32) error {
 
 // BeforeServiceDeleted implements servicestypes.ServicesHooks
 func (h ServicesHooks) BeforeServiceDeleted(ctx context.Context, serviceID uint32) error {
-	return h.k.BeforeDelegationTargetRemoved(ctx, restakingtypes.DELEGATION_TYPE_SERVICE, serviceID)
+	return h.k.BeforeServiceDeleted(ctx, serviceID)
 }
 
 // AfterServiceAccreditationModified implements servicestypes.ServicesHooks