From 2c1105398be27f27d28136e72bdf0e0835728607 Mon Sep 17 00:00:00 2001 From: milad Date: Tue, 24 Feb 2026 12:15:20 +0330 Subject: [PATCH 01/34] temp data structurues for pse --- x/pse/keeper/delegation.go | 7 +------ x/pse/keeper/distribute_test.go | 4 ++-- x/pse/keeper/genesis.go | 8 ++++++++ x/pse/types/distribution.pb.go | 9 +++++---- 4 files changed, 16 insertions(+), 12 deletions(-) diff --git a/x/pse/keeper/delegation.go b/x/pse/keeper/delegation.go index 07acd176..b3da8459 100644 --- a/x/pse/keeper/delegation.go +++ b/x/pse/keeper/delegation.go @@ -56,12 +56,7 @@ func (k Keeper) GetDelegatorScore( } // SetDelegatorScore sets the score for a delegator. -func (k Keeper) SetDelegatorScore( - ctx context.Context, - distributionID uint64, - delAddr sdk.AccAddress, - score sdkmath.Int, -) error { +func (k Keeper) SetDelegatorScore(ctx context.Context, distributionID uint64, delAddr sdk.AccAddress, score sdkmath.Int) error { key := collections.Join(distributionID, delAddr) return k.AccountScoreSnapshot.Set(ctx, key, score) } diff --git a/x/pse/keeper/distribute_test.go b/x/pse/keeper/distribute_test.go index 675ba11b..16e6f32e 100644 --- a/x/pse/keeper/distribute_test.go +++ b/x/pse/keeper/distribute_test.go @@ -33,8 +33,8 @@ func TestKeeper_Distribute(t *testing.T) { func(r *runEnv) { distributeAction(r, sdkmath.NewInt(1000)) }, func(r *runEnv) { assertDistributionAction(r, map[*sdk.AccAddress]sdkmath.Int{ - &r.delegators[0]: sdkmath.NewInt(1_100_366), // + 1000 * 1.1 / 3 - &r.delegators[1]: sdkmath.NewInt(900_299), // + 1000 * 0.9 / 3 + &r.delegators[0]: sdkmath.NewInt(1_100_550), // + 1000 * 1.1 / 3 + &r.delegators[1]: sdkmath.NewInt(900_450), // + 1000 * 0.9 / 3 }) }, func(r *runEnv) { assertScoreResetAction(r) }, diff --git a/x/pse/keeper/genesis.go b/x/pse/keeper/genesis.go index 9147ca6b..212e4751 100644 --- a/x/pse/keeper/genesis.go +++ b/x/pse/keeper/genesis.go @@ -28,6 +28,12 @@ func (k Keeper) InitGenesis(ctx context.Context, genState types.GenesisState) er } } + // TODO revise this logic for distribution id and genesis state + var currentDistributionID uint64 + if len(genState.ScheduledDistributions) > 0 { + currentDistributionID = genState.ScheduledDistributions[0].Timestamp + } + // Populate delegation time entries from genesis state for _, delegationTimeEntryExported := range genState.DelegationTimeEntries { valAddr, err := k.valAddressCodec.StringToBytes(delegationTimeEntryExported.ValidatorAddress) @@ -96,6 +102,7 @@ func (k Keeper) ExportGenesis(ctx context.Context) (*types.GenesisState, error) if err != nil { return false, err } + // TODO revise this logic for distribution id and genesis state delegationTimeEntriesExported = append(delegationTimeEntriesExported, types.DelegationTimeEntryExport{ DistributionID: key.K1(), ValidatorAddress: valAddr, @@ -117,6 +124,7 @@ func (k Keeper) ExportGenesis(ctx context.Context) (*types.GenesisState, error) if err != nil { return false, err } + // TODO revise this logic for distribution id and genesis state genesis.AccountScores = append(genesis.AccountScores, types.AccountScore{ DistributionID: key.K1(), Address: addr, diff --git a/x/pse/types/distribution.pb.go b/x/pse/types/distribution.pb.go index ef56433c..f1055c0a 100644 --- a/x/pse/types/distribution.pb.go +++ b/x/pse/types/distribution.pb.go @@ -4,14 +4,15 @@ package types import ( - cosmossdk_io_math "cosmossdk.io/math" fmt "fmt" - _ "github.com/cosmos/cosmos-proto" - _ "github.com/cosmos/gogoproto/gogoproto" - proto "github.com/cosmos/gogoproto/proto" io "io" math "math" math_bits "math/bits" + + cosmossdk_io_math "cosmossdk.io/math" + _ "github.com/cosmos/cosmos-proto" + _ "github.com/cosmos/gogoproto/gogoproto" + proto "github.com/cosmos/gogoproto/proto" ) // Reference imports to suppress errors if they are not otherwise used. From cf9684b6c356e6050d5c5f42ba397c0bd9669cd3 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Thu, 26 Feb 2026 12:14:33 +0300 Subject: [PATCH 02/34] add TotalScore and OngoingDistribution storages to keeper --- x/pse/keeper/keeper.go | 17 ++++++++++++++++- x/pse/types/key.go | 4 +++- 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/x/pse/keeper/keeper.go b/x/pse/keeper/keeper.go index a66c1d20..53363676 100644 --- a/x/pse/keeper/keeper.go +++ b/x/pse/keeper/keeper.go @@ -37,7 +37,9 @@ type Keeper struct { types.DelegationTimeEntry, ] AccountScoreSnapshot collections.Map[collections.Pair[uint64, sdk.AccAddress], sdkmath.Int] - AllocationSchedule collections.Map[uint64, types.ScheduledDistribution] // Map: id -> ScheduledDistribution + AllocationSchedule collections.Map[uint64, types.ScheduledDistribution] // Map: ID -> ScheduledDistribution + TotalScore collections.Map[uint64, sdkmath.Int] // Map: ID -> total accumulated score + OngoingDistribution collections.Item[types.ScheduledDistribution] // Currently processing distribution DistributionDisabled collections.Item[bool] } @@ -92,6 +94,19 @@ func NewKeeper( collections.Uint64Key, codec.CollValue[types.ScheduledDistribution](cdc), ), + TotalScore: collections.NewMap( + sb, + types.TotalScoreKey, + "total_score", + collections.Uint64Key, + sdk.IntValue, + ), + OngoingDistribution: collections.NewItem( + sb, + types.OngoingDistributionKey, + "ongoing_distribution", + codec.CollValue[types.ScheduledDistribution](cdc), + ), DistributionDisabled: collections.NewItem( sb, types.DistributionDisabledKey, diff --git a/x/pse/types/key.go b/x/pse/types/key.go index be6e50fe..91734ab5 100644 --- a/x/pse/types/key.go +++ b/x/pse/types/key.go @@ -15,6 +15,8 @@ var ( ParamsKey = collections.NewPrefix(0) StakingTimeKey = collections.NewPrefix(1) AccountScoreKey = collections.NewPrefix(2) - AllocationScheduleKey = collections.NewPrefix(3) // Map: timestamp -> ScheduledDistribution + AllocationScheduleKey = collections.NewPrefix(3) // Map: ID -> ScheduledDistribution DistributionDisabledKey = collections.NewPrefix(4) + TotalScoreKey = collections.NewPrefix(5) // Map: ID -> total accumulated score + OngoingDistributionKey = collections.NewPrefix(6) // Item: currently processing ScheduledDistribution ) From 43391895c91a5892f2fe4e919ed348ab1cbab4de Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Thu, 26 Feb 2026 12:15:38 +0300 Subject: [PATCH 03/34] add addToScore helpers for atomic score updates --- x/pse/keeper/delegation.go | 30 +++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/x/pse/keeper/delegation.go b/x/pse/keeper/delegation.go index b3da8459..5cf58749 100644 --- a/x/pse/keeper/delegation.go +++ b/x/pse/keeper/delegation.go @@ -67,6 +67,34 @@ func (k Keeper) RemoveDelegatorScore(ctx context.Context, distributionID uint64, return k.AccountScoreSnapshot.Remove(ctx, key) } +// addToScore atomically adds a score value to a delegator's score snapshot. +func (k Keeper) addToScore(ctx context.Context, distributionID uint64, delAddr sdk.AccAddress, score sdkmath.Int) error { + if score.IsZero() { + return nil + } + lastScore, err := k.GetDelegatorScore(ctx, distributionID, delAddr) + if errors.Is(err, collections.ErrNotFound) { + lastScore = sdkmath.NewInt(0) + } else if err != nil { + return err + } + return k.SetDelegatorScore(ctx, distributionID, delAddr, lastScore.Add(score)) +} + +// addToTotalScore atomically adds a score value to a distribution's total score. +func (k Keeper) addToTotalScore(ctx context.Context, distributionID uint64, score sdkmath.Int) error { + if score.IsZero() { + return nil + } + current, err := k.TotalScore.Get(ctx, distributionID) + if errors.Is(err, collections.ErrNotFound) { + current = sdkmath.NewInt(0) + } else if err != nil { + return err + } + return k.TotalScore.Set(ctx, distributionID, current.Add(score)) +} + // CalculateDelegatorScore calculates the current total score for a delegator. // This includes both the accumulated score snapshot (from previous periods) // and the current period score calculated on-demand from active delegations. @@ -77,7 +105,7 @@ func (k Keeper) CalculateDelegatorScore(ctx context.Context, delAddr sdk.AccAddr if err != nil { return sdkmath.Int{}, err } - distributionID := distribution.ID // TODO update to handle distribution ID properly. + distributionID := distribution.ID // Start with the accumulated score from the snapshot (previous periods) accumulatedScore, err := k.GetDelegatorScore(ctx, distributionID, delAddr) From bb034c261e6269e79e24d93f390bc0ced1aa17ad Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Thu, 26 Feb 2026 12:22:42 +0300 Subject: [PATCH 04/34] rewrite hooks with 3-scenario logic for multi-block distribution --- x/pse/keeper/hooks.go | 243 +++++++++++++++++++++++++++++++----------- 1 file changed, 182 insertions(+), 61 deletions(-) diff --git a/x/pse/keeper/hooks.go b/x/pse/keeper/hooks.go index 5c00b8ec..1559850d 100644 --- a/x/pse/keeper/hooks.go +++ b/x/pse/keeper/hooks.go @@ -24,35 +24,57 @@ func (k Keeper) Hooks() Hooks { return Hooks{k} } +// getOngoingDistribution returns the ongoing distribution if one exists, or nil if not. +func (k Keeper) getOngoingDistribution(ctx context.Context) (*types.ScheduledDistribution, error) { + ongoing, err := k.OngoingDistribution.Get(ctx) + if errors.Is(err, collections.ErrNotFound) { + return nil, nil + } + if err != nil { + return nil, err + } + return &ongoing, nil +} + +// getCurrentDistributionID returns the distribution ID that new entries should be written to. +// If an ongoing distribution exists (ID=N is being processed), returns N+1. +// Otherwise returns the next scheduled distribution's ID. +// Returns 0 if no distribution is scheduled and none is ongoing. +func (k Keeper) getCurrentDistributionID(ctx context.Context) (uint64, error) { + ongoing, err := k.getOngoingDistribution(ctx) + if err != nil { + return 0, err + } + if ongoing != nil { + return ongoing.ID + 1, nil + } + + distribution, _, err := k.PeekNextAllocationSchedule(ctx) + if err != nil { + return 0, err + } + return distribution.ID, nil +} + // AfterDelegationModified implements the staking hooks interface. +// Handles 3 scenarios based on where the delegator's entry exists: +// - Scenario 1: Entry in prevID (ongoing distribution in progress) — +// - Scenario 2: Entry in currentID — normal score calculation. +// - Scenario 3: No entry — create new entry, no score. func (h Hooks) AfterDelegationModified(ctx context.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress) error { delegation, err := h.k.stakingKeeper.GetDelegation(ctx, delAddr, valAddr) if err != nil { return err } - // TODO handle empty distribution schedule in all places. - distribution, _, err := h.k.PeekNextAllocationSchedule(ctx) + currentID, err := h.k.getCurrentDistributionID(ctx) if err != nil { return err } - // TODO update to handle distribution ID properly. - // We should update the logic to find the active distribution (probably via a new store called OngoingDistribution) - // and check for period splits - distributionID := distribution.ID - - blockTimeUnixSeconds := sdk.UnwrapSDKContext(ctx).BlockTime().Unix() - delegationTimeEntry, err := h.k.GetDelegationTimeEntry(ctx, distributionID, valAddr, delAddr) - if errors.Is(err, collections.ErrNotFound) { - delegationTimeEntry = types.DelegationTimeEntry{ - LastChangedUnixSec: blockTimeUnixSeconds, - Shares: delegation.Shares, - } - } else if err != nil { - return err + if currentID == 0 { + return nil } - // Stop score addition for excluded addresses isExcluded, err := h.k.IsExcludedAddress(ctx, delAddr) if err != nil { return err @@ -61,48 +83,92 @@ func (h Hooks) AfterDelegationModified(ctx context.Context, delAddr sdk.AccAddre return nil } - // Only update AccountScoreSnapshot for non-excluded addresses - lastScore, err := h.k.GetDelegatorScore(ctx, distributionID, delAddr) - if errors.Is(err, collections.ErrNotFound) { - lastScore = sdkmath.NewInt(0) - } else if err != nil { - return err - } + blockTime := sdk.UnwrapSDKContext(ctx).BlockTime().Unix() - addedScore, err := calculateAddedScore(ctx, h.k, valAddr, delegationTimeEntry) + // Scenario 1: Entry exists in previous distribution (ongoing distribution in progress). + // Split score at distribution timestamp, move entry to currentID. + ongoing, err := h.k.getOngoingDistribution(ctx) if err != nil { return err } - newScore := lastScore.Add(addedScore) + if ongoing != nil { + prevID := ongoing.ID + prevEntry, err := h.k.GetDelegationTimeEntry(ctx, prevID, valAddr, delAddr) + if err == nil { + distTimestamp := int64(ongoing.Timestamp) + + // Score for previous period: lastChanged -> distribution timestamp + prevScore, err := calculateScoreAtTimestamp(ctx, h.k, valAddr, prevEntry, distTimestamp) + if err != nil { + return err + } + if err := h.k.addToScore(ctx, prevID, delAddr, prevScore); err != nil { + return err + } + + // Score for current period: distribution timestamp -> now (old shares still active) + currentPeriodEntry := types.DelegationTimeEntry{ + LastChangedUnixSec: distTimestamp, + Shares: prevEntry.Shares, + } + currentScore, err := calculateScoreAtTimestamp(ctx, h.k, valAddr, currentPeriodEntry, blockTime) + if err != nil { + return err + } + if err := h.k.addToScore(ctx, currentID, delAddr, currentScore); err != nil { + return err + } + + // Delete from prevID, create in currentID with new shares + if err := h.k.RemoveDelegationTimeEntry(ctx, prevID, valAddr, delAddr); err != nil { + return err + } + return h.k.SetDelegationTimeEntry(ctx, currentID, valAddr, delAddr, types.DelegationTimeEntry{ + LastChangedUnixSec: blockTime, + Shares: delegation.Shares, + }) + } + if !errors.Is(err, collections.ErrNotFound) { + return err + } + } - // Update DelegationTimeEntry for non-excluded addresses - if err := h.k.SetDelegationTimeEntry(ctx, distributionID, valAddr, delAddr, types.DelegationTimeEntry{ - LastChangedUnixSec: blockTimeUnixSeconds, - Shares: delegation.Shares, - }); err != nil { + // Scenario 2: Entry exists in current distribution + currentEntry, err := h.k.GetDelegationTimeEntry(ctx, currentID, valAddr, delAddr) + if err == nil { + score, err := calculateAddedScore(ctx, h.k, valAddr, currentEntry) + if err != nil { + return err + } + if err := h.k.addToScore(ctx, currentID, delAddr, score); err != nil { + return err + } + return h.k.SetDelegationTimeEntry(ctx, currentID, valAddr, delAddr, types.DelegationTimeEntry{ + LastChangedUnixSec: blockTime, + Shares: delegation.Shares, + }) + } + if !errors.Is(err, collections.ErrNotFound) { return err } - return h.k.SetDelegatorScore(ctx, distributionID, delAddr, newScore) + // Scenario 3: No entry - create new in currentID (no score, duration = 0) + return h.k.SetDelegationTimeEntry(ctx, currentID, valAddr, delAddr, types.DelegationTimeEntry{ + LastChangedUnixSec: blockTime, + Shares: delegation.Shares, + }) } // BeforeDelegationRemoved implements the staking hooks interface. func (h Hooks) BeforeDelegationRemoved(ctx context.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress) error { - distribution, _, err := h.k.PeekNextAllocationSchedule(ctx) + currentID, err := h.k.getCurrentDistributionID(ctx) if err != nil { return err } - distributionID := distribution.ID // TODO update to handle distribution ID properly. - - delegationTimeEntry, err := h.k.GetDelegationTimeEntry(ctx, distributionID, valAddr, delAddr) - if err != nil { - if errors.Is(err, collections.ErrNotFound) { - return nil - } - return err + if currentID == 0 { + return nil } - // Stop score addition for excluded addresses isExcluded, err := h.k.IsExcludedAddress(ctx, delAddr) if err != nil { return err @@ -111,44 +177,99 @@ func (h Hooks) BeforeDelegationRemoved(ctx context.Context, delAddr sdk.AccAddre return nil } - // Only update AccountScoreSnapshot for non-excluded addresses - lastScore, err := h.k.GetDelegatorScore(ctx, distributionID, delAddr) - if errors.Is(err, collections.ErrNotFound) { - lastScore = sdkmath.NewInt(0) - } else if err != nil { - return err - } + blockTime := sdk.UnwrapSDKContext(ctx).BlockTime().Unix() - addedScore, err := calculateAddedScore(ctx, h.k, valAddr, delegationTimeEntry) + // Scenario 1: Entry exists in previous distribution (ongoing) + ongoing, err := h.k.getOngoingDistribution(ctx) if err != nil { return err } - newScore := lastScore.Add(addedScore) + if ongoing != nil { + prevID := ongoing.ID + prevEntry, err := h.k.GetDelegationTimeEntry(ctx, prevID, valAddr, delAddr) + if err == nil { + distTimestamp := int64(ongoing.Timestamp) + + // Score for previous period: lastChanged -> distribution timestamp + prevScore, err := calculateScoreAtTimestamp(ctx, h.k, valAddr, prevEntry, distTimestamp) + if err != nil { + return err + } + if err := h.k.addToScore(ctx, prevID, delAddr, prevScore); err != nil { + return err + } + + // Score for current period: distribution timestamp -> now + currentPeriodEntry := types.DelegationTimeEntry{ + LastChangedUnixSec: distTimestamp, + Shares: prevEntry.Shares, + } + currentScore, err := calculateScoreAtTimestamp(ctx, h.k, valAddr, currentPeriodEntry, blockTime) + if err != nil { + return err + } + if err := h.k.addToScore(ctx, currentID, delAddr, currentScore); err != nil { + return err + } + + // Delete from prevID (delegation removed) + return h.k.RemoveDelegationTimeEntry(ctx, prevID, valAddr, delAddr) + } + if !errors.Is(err, collections.ErrNotFound) { + return err + } + } - // Remove DelegationTimeEntry for non-excluded addresses - if err := h.k.RemoveDelegationTimeEntry(ctx, distributionID, valAddr, delAddr); err != nil { + // Scenario 2: Entry exists in current distribution + currentEntry, err := h.k.GetDelegationTimeEntry(ctx, currentID, valAddr, delAddr) + if err == nil { + score, err := calculateAddedScore(ctx, h.k, valAddr, currentEntry) + if err != nil { + return err + } + if err := h.k.addToScore(ctx, currentID, delAddr, score); err != nil { + return err + } + return h.k.RemoveDelegationTimeEntry(ctx, currentID, valAddr, delAddr) + } + if !errors.Is(err, collections.ErrNotFound) { return err } - return h.k.SetDelegatorScore(ctx, distributionID, delAddr, newScore) + // Scenario 3: No entry + return nil } -func calculateAddedScore( +// calculateScoreAtTimestamp calculates the score for a delegation entry up to a specific timestamp. +// score = tokens × (atTimestamp - lastChanged) +func calculateScoreAtTimestamp( ctx context.Context, keeper Keeper, valAddr sdk.ValAddress, - delegationTimeEntry types.DelegationTimeEntry, + entry types.DelegationTimeEntry, + atTimestamp int64, ) (sdkmath.Int, error) { val, err := keeper.stakingKeeper.GetValidator(ctx, valAddr) if err != nil { return sdkmath.NewInt(0), err } + duration := atTimestamp - entry.LastChangedUnixSec + if duration <= 0 { + return sdkmath.NewInt(0), nil + } + tokens := val.TokensFromShares(entry.Shares).TruncateInt() + return tokens.MulRaw(duration), nil +} - blockTimeUnixSeconds := sdk.UnwrapSDKContext(ctx).BlockTime().Unix() - delegationDuration := blockTimeUnixSeconds - delegationTimeEntry.LastChangedUnixSec - previousDelegatedTokens := val.TokensFromShares(delegationTimeEntry.Shares).TruncateInt() - delegationScore := previousDelegatedTokens.MulRaw(delegationDuration) - return delegationScore, nil +// calculateAddedScore calculates the score for a delegation entry up to the current block time. +func calculateAddedScore( + ctx context.Context, + keeper Keeper, + valAddr sdk.ValAddress, + delegationTimeEntry types.DelegationTimeEntry, +) (sdkmath.Int, error) { + blockTime := sdk.UnwrapSDKContext(ctx).BlockTime().Unix() + return calculateScoreAtTimestamp(ctx, keeper, valAddr, delegationTimeEntry, blockTime) } // BeforeValidatorSlashed implements the staking hooks interface. From 990fa47288f058704f20b1e7420b767a53235a6e Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Thu, 26 Feb 2026 14:20:40 +0300 Subject: [PATCH 05/34] remove redundant addToTotalScore func --- x/pse/keeper/delegation.go | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/x/pse/keeper/delegation.go b/x/pse/keeper/delegation.go index 5cf58749..141f22f1 100644 --- a/x/pse/keeper/delegation.go +++ b/x/pse/keeper/delegation.go @@ -81,20 +81,6 @@ func (k Keeper) addToScore(ctx context.Context, distributionID uint64, delAddr s return k.SetDelegatorScore(ctx, distributionID, delAddr, lastScore.Add(score)) } -// addToTotalScore atomically adds a score value to a distribution's total score. -func (k Keeper) addToTotalScore(ctx context.Context, distributionID uint64, score sdkmath.Int) error { - if score.IsZero() { - return nil - } - current, err := k.TotalScore.Get(ctx, distributionID) - if errors.Is(err, collections.ErrNotFound) { - current = sdkmath.NewInt(0) - } else if err != nil { - return err - } - return k.TotalScore.Set(ctx, distributionID, current.Add(score)) -} - // CalculateDelegatorScore calculates the current total score for a delegator. // This includes both the accumulated score snapshot (from previous periods) // and the current period score calculated on-demand from active delegations. From df904d0ea8292184542f407bef26f7200884b85c Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Fri, 27 Feb 2026 11:06:22 +0300 Subject: [PATCH 06/34] add phase-1 batched score conversion --- x/pse/keeper/distribute.go | 111 +++++++++++++++++++++++++++++++++++++ x/pse/keeper/keeper.go | 8 +++ x/pse/types/key.go | 1 + 3 files changed, 120 insertions(+) diff --git a/x/pse/keeper/distribute.go b/x/pse/keeper/distribute.go index 091f1a51..ba3cefb1 100644 --- a/x/pse/keeper/distribute.go +++ b/x/pse/keeper/distribute.go @@ -11,6 +11,117 @@ import ( "github.com/tokenize-x/tx-chain/v7/x/pse/types" ) +// defaultBatchSize is the number of entries processed per EndBlock during multi-block distribution. +const defaultBatchSize = 100 // TODO: make configurable + +// ProcessPhase1ScoreConversion processes a batch of DelegationTimeEntries from the ongoing distribution (prevID), +// converting each entry into a score snapshot and migrating it to currentID (prevID+1). +// +// For each entry in the batch: +// 1. Calculate score from lastChanged to distribution timestamp → addToScore(prevID) +// 2. Create new entry in currentID with same shares, lastChanged = distTimestamp +// 3. Remove entry from prevID +// +// Returns true when all prevID entries have been processed and TotalScore is computed. +func (k Keeper) ProcessPhase1ScoreConversion(ctx context.Context, ongoing types.ScheduledDistribution) (bool, error) { + prevID := ongoing.ID + currentID := ongoing.ID + 1 + distTimestamp := int64(ongoing.Timestamp) + + // Collect a batch of entries from prevID. + iter, err := k.DelegationTimeEntries.Iterate( + ctx, + collections.NewPrefixedTripleRange[uint64, sdk.AccAddress, sdk.ValAddress](prevID), + ) + if err != nil { + return false, err + } + + type entryKV struct { + delAddr sdk.AccAddress + valAddr sdk.ValAddress + entry types.DelegationTimeEntry + } + + batch := make([]entryKV, 0, defaultBatchSize) + for ; iter.Valid() && len(batch) < defaultBatchSize; iter.Next() { + kv, err := iter.KeyValue() + if err != nil { + iter.Close() + return false, err + } + batch = append(batch, entryKV{ + delAddr: kv.Key.K2(), + valAddr: kv.Key.K3(), + entry: kv.Value, + }) + } + iter.Close() + + // Compute TotalScore from all accumulated snapshots. + if len(batch) == 0 { + if err := k.computeTotalScore(ctx, prevID); err != nil { + return false, err + } + return true, nil + } + + for _, item := range batch { + isExcluded, err := k.IsExcludedAddress(ctx, item.delAddr) + if err != nil { + return false, err + } + + if !isExcluded { + score, err := calculateScoreAtTimestamp(ctx, k, item.valAddr, item.entry, distTimestamp) + if err != nil { + return false, err + } + if err := k.addToScore(ctx, prevID, item.delAddr, score); err != nil { + return false, err + } + } + + // Migrate entry to currentID with same shares, reset lastChanged to distribution timestamp. + if err := k.SetDelegationTimeEntry(ctx, currentID, item.valAddr, item.delAddr, types.DelegationTimeEntry{ + LastChangedUnixSec: distTimestamp, + Shares: item.entry.Shares, + }); err != nil { + return false, err + } + + // Remove from prevID. + if err := k.RemoveDelegationTimeEntry(ctx, prevID, item.valAddr, item.delAddr); err != nil { + return false, err + } + } + + return false, nil +} + +// computeTotalScore sums all AccountScoreSnapshot entries for a distribution and stores the result in TotalScore. +func (k Keeper) computeTotalScore(ctx context.Context, distributionID uint64) error { + iter, err := k.AccountScoreSnapshot.Iterate( + ctx, + collections.NewPrefixedPairRange[uint64, sdk.AccAddress](distributionID), + ) + if err != nil { + return err + } + defer iter.Close() + + totalScore := sdkmath.NewInt(0) + for ; iter.Valid(); iter.Next() { + kv, err := iter.KeyValue() + if err != nil { + return err + } + totalScore = totalScore.Add(kv.Value) + } + + return k.TotalScore.Set(ctx, distributionID, totalScore) +} + // DistributeCommunityPSE distributes the total community PSE amount to all delegators based on their score. func (k Keeper) DistributeCommunityPSE( ctx context.Context, diff --git a/x/pse/keeper/keeper.go b/x/pse/keeper/keeper.go index 53363676..4d07f0ef 100644 --- a/x/pse/keeper/keeper.go +++ b/x/pse/keeper/keeper.go @@ -40,6 +40,7 @@ type Keeper struct { AllocationSchedule collections.Map[uint64, types.ScheduledDistribution] // Map: ID -> ScheduledDistribution TotalScore collections.Map[uint64, sdkmath.Int] // Map: ID -> total accumulated score OngoingDistribution collections.Item[types.ScheduledDistribution] // Currently processing distribution + DistributedAmount collections.Map[uint64, sdkmath.Int] // Map: ID -> cumulative distributed amount DistributionDisabled collections.Item[bool] } @@ -107,6 +108,13 @@ func NewKeeper( "ongoing_distribution", codec.CollValue[types.ScheduledDistribution](cdc), ), + DistributedAmount: collections.NewMap( + sb, + types.DistributedAmountKey, + "distributed_amount", + collections.Uint64Key, + sdk.IntValue, + ), DistributionDisabled: collections.NewItem( sb, types.DistributionDisabledKey, diff --git a/x/pse/types/key.go b/x/pse/types/key.go index 91734ab5..040e8de3 100644 --- a/x/pse/types/key.go +++ b/x/pse/types/key.go @@ -19,4 +19,5 @@ var ( DistributionDisabledKey = collections.NewPrefix(4) TotalScoreKey = collections.NewPrefix(5) // Map: ID -> total accumulated score OngoingDistributionKey = collections.NewPrefix(6) // Item: currently processing ScheduledDistribution + DistributedAmountKey = collections.NewPrefix(7) // Map: ID -> cumulative distributed amount ) From fac88cf5afb8285c1ba23b6d243fa1e278db5da9 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Fri, 27 Feb 2026 11:58:52 +0300 Subject: [PATCH 07/34] add phase-2 batched token distribution --- x/pse/keeper/distribute.go | 165 +++++++++++++++++++++++++++++++++++++ 1 file changed, 165 insertions(+) diff --git a/x/pse/keeper/distribute.go b/x/pse/keeper/distribute.go index ba3cefb1..33e0cdff 100644 --- a/x/pse/keeper/distribute.go +++ b/x/pse/keeper/distribute.go @@ -2,6 +2,7 @@ package keeper import ( "context" + "errors" "cosmossdk.io/collections" sdkmath "cosmossdk.io/math" @@ -122,6 +123,170 @@ func (k Keeper) computeTotalScore(ctx context.Context, distributionID uint64) er return k.TotalScore.Set(ctx, distributionID, totalScore) } +// ProcessPhase2TokenDistribution distributes tokens to delegators in batches based on their computed scores. +// Uses TotalScore[prevID] for proportion calculation and iterates AccountScoreSnapshot[prevID]. +// +// For each delegator in the batch: +// 1. Compute share: userAmount = totalPSEAmount × score / totalScore +// 2. Distribute via distributeToDelegator (send tokens + auto-delegate) +// 3. Track cumulative distributed amount +// 4. Remove the processed snapshot entry +// +// When all delegators have been processed, sends leftover (rounding errors + undelegated users) to the community pool. +// Returns true when distribution is complete and all state has been cleaned up. +func (k Keeper) ProcessPhase2TokenDistribution(ctx context.Context, ongoing types.ScheduledDistribution, bondDenom string) (bool, error) { + prevID := ongoing.ID + totalPSEAmount := getCommunityAllocationAmount(ongoing) + + totalScore, err := k.TotalScore.Get(ctx, prevID) + if err != nil { + return false, err + } + + // No score or no amount: send everything to community pool and clean up. + if !totalScore.IsPositive() || totalPSEAmount.IsZero() { + if totalPSEAmount.IsPositive() { + if err := k.sendLeftoverToCommunityPool(ctx, totalPSEAmount, bondDenom); err != nil { + return false, err + } + } + return true, k.cleanupDistribution(ctx, prevID) + } + + // Collect a batch of score snapshots. + iter, err := k.AccountScoreSnapshot.Iterate( + ctx, + collections.NewPrefixedPairRange[uint64, sdk.AccAddress](prevID), + ) + if err != nil { + return false, err + } + + type scoreEntry struct { + delAddr sdk.AccAddress + score sdkmath.Int + } + + batch := make([]scoreEntry, 0, defaultBatchSize) + for ; iter.Valid() && len(batch) < defaultBatchSize; iter.Next() { + kv, err := iter.KeyValue() + if err != nil { + iter.Close() + return false, err + } + batch = append(batch, scoreEntry{ + delAddr: kv.Key.K2(), + score: kv.Value, + }) + } + iter.Close() + + // Only triggered when all distributions of this round are completed. + // Send leftover to community pool and clean up. + if len(batch) == 0 { + distributedSoFar, err := k.getDistributedAmount(ctx, prevID) + if err != nil { + return false, err + } + leftover := totalPSEAmount.Sub(distributedSoFar) + if leftover.IsPositive() { + if err := k.sendLeftoverToCommunityPool(ctx, leftover, bondDenom); err != nil { + return false, err + } + } + return true, k.cleanupDistribution(ctx, prevID) + } + + sdkCtx := sdk.UnwrapSDKContext(ctx) + + // Distribute rewards to each delegator in the batch proportional to their score. + for _, item := range batch { + userAmount := totalPSEAmount.Mul(item.score).Quo(totalScore) + distributedAmount, err := k.distributeToDelegator(ctx, item.delAddr, userAmount, bondDenom) + if err != nil { + return false, err + } + + if err := k.addToDistributedAmount(ctx, prevID, distributedAmount); err != nil { + return false, err + } + + if err := sdkCtx.EventManager().EmitTypedEvent(&types.EventCommunityDistributed{ + DelegatorAddress: item.delAddr.String(), + Score: item.score, + TotalPseScore: totalScore, + Amount: userAmount, + ScheduledAt: ongoing.Timestamp, + }); err != nil { + sdkCtx.Logger().Error("failed to emit community distributed event", "error", err) + } + + // Remove processed snapshot. + if err := k.RemoveDelegatorScore(ctx, prevID, item.delAddr); err != nil { + return false, err + } + } + + return false, nil +} + +// getCommunityAllocationAmount extracts the community clearing account allocation from a distribution. +func getCommunityAllocationAmount(dist types.ScheduledDistribution) sdkmath.Int { + for _, alloc := range dist.Allocations { + if alloc.ClearingAccount == types.ClearingAccountCommunity { + return alloc.Amount + } + } + return sdkmath.NewInt(0) +} + +// sendLeftoverToCommunityPool sends remaining undistributed tokens to the community pool. +func (k Keeper) sendLeftoverToCommunityPool(ctx context.Context, amount sdkmath.Int, bondDenom string) error { + pseModuleAddress := k.accountKeeper.GetModuleAddress(types.ClearingAccountCommunity) + return k.distributionKeeper.FundCommunityPool(ctx, sdk.NewCoins(sdk.NewCoin(bondDenom, amount)), pseModuleAddress) +} + +// cleanupDistribution removes all state associated with a completed distribution. +func (k Keeper) cleanupDistribution(ctx context.Context, distributionID uint64) error { + if err := k.AccountScoreSnapshot.Clear( + ctx, + collections.NewPrefixedPairRange[uint64, sdk.AccAddress](distributionID), + ); err != nil { + return err + } + if err := k.TotalScore.Remove(ctx, distributionID); err != nil { + return err + } + if err := k.DistributedAmount.Remove(ctx, distributionID); err != nil { + return err + } + if err := k.AllocationSchedule.Remove(ctx, distributionID); err != nil { + return err + } + return k.OngoingDistribution.Remove(ctx) +} + +// getDistributedAmount returns the cumulative distributed amount for a distribution, or zero if not set. +func (k Keeper) getDistributedAmount(ctx context.Context, distributionID uint64) (sdkmath.Int, error) { + amount, err := k.DistributedAmount.Get(ctx, distributionID) + if errors.Is(err, collections.ErrNotFound) { + return sdkmath.NewInt(0), nil + } + return amount, err +} + +// addToDistributedAmount atomically adds to the cumulative distributed amount. +func (k Keeper) addToDistributedAmount(ctx context.Context, distributionID uint64, amount sdkmath.Int) error { + if amount.IsZero() { + return nil + } + current, err := k.getDistributedAmount(ctx, distributionID) + if err != nil { + return err + } + return k.DistributedAmount.Set(ctx, distributionID, current.Add(amount)) +} + // DistributeCommunityPSE distributes the total community PSE amount to all delegators based on their score. func (k Keeper) DistributeCommunityPSE( ctx context.Context, From 22e966f6e3d1a0cbe82fc30da120cbd658c02afd Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Fri, 27 Feb 2026 12:25:02 +0300 Subject: [PATCH 08/34] rewrite EndBlocker with multi-block phase routing --- x/pse/keeper/distribution.go | 182 +++++++++++++++++++++++++++++++---- 1 file changed, 164 insertions(+), 18 deletions(-) diff --git a/x/pse/keeper/distribution.go b/x/pse/keeper/distribution.go index 19e47ca2..2a9ddec8 100644 --- a/x/pse/keeper/distribution.go +++ b/x/pse/keeper/distribution.go @@ -2,7 +2,9 @@ package keeper import ( "context" + "errors" + "cosmossdk.io/collections" errorsmod "cosmossdk.io/errors" sdkmath "cosmossdk.io/math" sdk "github.com/cosmos/cosmos-sdk/types" @@ -10,52 +12,196 @@ import ( "github.com/tokenize-x/tx-chain/v7/x/pse/types" ) -// ProcessNextDistribution processes the next due distribution from the schedule. -// Checks the earliest scheduled distribution and processes it if the current block time has passed its timestamp. -// Only one distribution is processed per call. Should be called from EndBlock. +// ProcessNextDistribution is the EndBlock entry point for distribution processing. +// It either resumes an ongoing multi-block distribution or starts a new one if due. +// 1. If OngoingDistribution exists -> resume (Phase 1 or Phase 2) +// 2. If no ongoing -> peek schedule -> if due: +// a. Process non-community allocations immediately (single-block) +// b. If community allocation exists -> set OngoingDistribution (Phase 1 starts next block) +// c. Else, no community allocation, non-community distribution is already done, remove from AllocationSchedule func (k Keeper) ProcessNextDistribution(ctx context.Context) error { - sdkCtx := sdk.UnwrapSDKContext(ctx) + // Resume ongoing multi-block distribution if one is in progress. + ongoing, err := k.getOngoingDistribution(ctx) + if err != nil { + return err + } + if ongoing != nil { + return k.resumeOngoingDistribution(ctx, *ongoing) + } - // Peek at the next scheduled distribution + // No ongoing distribution — check if next scheduled distribution is due. scheduledDistribution, shouldProcess, err := k.PeekNextAllocationSchedule(ctx) if err != nil { return err } - - // Return early if schedule is empty or not ready to process if !shouldProcess { return nil } - timestamp := scheduledDistribution.Timestamp - - // Get bond denom from staking params - //nolint:contextcheck // this is correct context passing bondDenom, err := k.stakingKeeper.BondDenom(ctx) if err != nil { return err } - // Get params containing clearing account to recipient address mappings params, err := k.GetParams(ctx) if err != nil { return err } - // Process all allocations scheduled for this timestamp - if err := k.distributeAllocatedTokens( - ctx, timestamp, bondDenom, params.ClearingAccountMappings, scheduledDistribution, + // Process non-community allocations immediately (single-block). + if err := k.distributeNonCommunityAllocations( + ctx, scheduledDistribution, bondDenom, params.ClearingAccountMappings, ); err != nil { return err } - // Remove the completed distribution from the schedule + sdkCtx := sdk.UnwrapSDKContext(ctx) + + // If community allocation exists, start multi-block processing. + communityAmount := getCommunityAllocationAmount(scheduledDistribution) + if communityAmount.IsPositive() { + if err := k.OngoingDistribution.Set(ctx, scheduledDistribution); err != nil { + return err + } + sdkCtx.Logger().Info("started multi-block community distribution", + "distribution_id", scheduledDistribution.ID, + "timestamp", scheduledDistribution.Timestamp) + return nil + } + + // No community allocation — remove from schedule if err := k.AllocationSchedule.Remove(ctx, scheduledDistribution.ID); err != nil { return err } - sdkCtx.Logger().Info("processed and removed allocation from schedule", - "timestamp", timestamp) + sdkCtx.Logger().Info("processed single-block distribution", + "distribution_id", scheduledDistribution.ID, + "timestamp", scheduledDistribution.Timestamp) + + return nil +} + +// resumeOngoingDistribution continues a multi-block community distribution. +// Phase is determined by TotalScore existence: absent -> Phase 1, present -> Phase 2. +func (k Keeper) resumeOngoingDistribution(ctx context.Context, ongoing types.ScheduledDistribution) error { + sdkCtx := sdk.UnwrapSDKContext(ctx) + prevID := ongoing.ID + + // TotalScore absent -> Phase 1 (score conversion still in progress). + _, err := k.TotalScore.Get(ctx, prevID) + if errors.Is(err, collections.ErrNotFound) { + done, err := k.ProcessPhase1ScoreConversion(ctx, ongoing) + if err != nil { + return err + } + if done { + sdkCtx.Logger().Info("phase 1 complete, TotalScore computed", + "distribution_id", prevID) + } + return nil + } + if err != nil { + return err + } + + // TotalScore present -> Phase 2 (token distribution). + bondDenom, err := k.stakingKeeper.BondDenom(ctx) + if err != nil { + return err + } + done, err := k.ProcessPhase2TokenDistribution(ctx, ongoing, bondDenom) + if err != nil { + return err + } + if done { + sdkCtx.Logger().Info("multi-block community distribution complete", + "distribution_id", prevID) + } + return nil +} + +// distributeNonCommunityAllocations processes all non-community allocations in a single block. +func (k Keeper) distributeNonCommunityAllocations( + ctx context.Context, + scheduledDistribution types.ScheduledDistribution, + bondDenom string, + clearingAccountMappings []types.ClearingAccountMapping, +) error { + sdkCtx := sdk.UnwrapSDKContext(ctx) + + for _, allocation := range scheduledDistribution.Allocations { + if allocation.Amount.IsZero() { + continue + } + + // Community allocation handled separately via multi-block distribution. + if allocation.ClearingAccount == types.ClearingAccountCommunity { + continue + } + + var recipientAddrs []string + for _, mapping := range clearingAccountMappings { + if mapping.ClearingAccount == allocation.ClearingAccount { + recipientAddrs = mapping.RecipientAddresses + break + } + } + + numRecipients := sdkmath.NewInt(int64(len(recipientAddrs))) + if numRecipients.IsZero() { + return errorsmod.Wrapf( + types.ErrTransferFailed, + "no recipients found for clearing account '%s'", + allocation.ClearingAccount, + ) + } + amountPerRecipient := allocation.Amount.Quo(numRecipients) + remainder := allocation.Amount.Mod(numRecipients) + + for _, recipientAddr := range recipientAddrs { + recipient := sdk.MustAccAddressFromBech32(recipientAddr) + coinsToSend := sdk.NewCoins(sdk.NewCoin(bondDenom, amountPerRecipient)) + + if err := k.bankKeeper.SendCoinsFromModuleToAccount( + ctx, + allocation.ClearingAccount, + recipient, + coinsToSend, + ); err != nil { + return errorsmod.Wrapf( + types.ErrTransferFailed, + "failed to transfer from clearing account '%s' to recipient '%s': %v", + allocation.ClearingAccount, + recipientAddr, + err, + ) + } + } + + if !remainder.IsZero() { + clearingAccountAddr := k.accountKeeper.GetModuleAddress(allocation.ClearingAccount) + remainderCoins := sdk.NewCoins(sdk.NewCoin(bondDenom, remainder)) + if err := k.distributionKeeper.FundCommunityPool(ctx, remainderCoins, clearingAccountAddr); err != nil { + return errorsmod.Wrapf( + types.ErrTransferFailed, + "failed to send remainder to community pool from clearing account '%s': %v", + allocation.ClearingAccount, + err, + ) + } + } + + if err := sdkCtx.EventManager().EmitTypedEvent(&types.EventAllocationDistributed{ + ClearingAccount: allocation.ClearingAccount, + RecipientAddresses: recipientAddrs, + AmountPerRecipient: amountPerRecipient, + CommunityPoolAmount: remainder, + ScheduledAt: scheduledDistribution.Timestamp, + TotalAmount: allocation.Amount, + }); err != nil { + sdkCtx.Logger().Error("failed to emit allocation completed event", "error", err) + } + } return nil } From 6e6aee12004027dabe4f4222f106f739ab3d0079 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Fri, 27 Feb 2026 12:56:09 +0300 Subject: [PATCH 09/34] reject schedule update during ongoing distribution --- x/pse/keeper/distribution.go | 9 +++++++++ x/pse/types/errors.go | 3 +++ 2 files changed, 12 insertions(+) diff --git a/x/pse/keeper/distribution.go b/x/pse/keeper/distribution.go index 2a9ddec8..8c0224a1 100644 --- a/x/pse/keeper/distribution.go +++ b/x/pse/keeper/distribution.go @@ -412,6 +412,15 @@ func (k Keeper) UpdateDistributionSchedule( return errorsmod.Wrapf(types.ErrInvalidAuthority, "expected %s, got %s", k.authority, authority) } + // Reject if a multi-block distribution is in progress. + ongoing, err := k.getOngoingDistribution(ctx) + if err != nil { + return err + } + if ongoing != nil { + return errorsmod.Wrapf(types.ErrOngoingDistribution, "cannot update schedule while distribution %d is in progress", ongoing.ID) + } + // Validate minimum gap between distributions params, err := k.GetParams(ctx) if err != nil { diff --git a/x/pse/types/errors.go b/x/pse/types/errors.go index eff36811..d279f870 100644 --- a/x/pse/types/errors.go +++ b/x/pse/types/errors.go @@ -22,4 +22,7 @@ var ( // ErrInvalidParam is returned when a parameter is invalid. ErrInvalidParam = sdkerrors.Register(ModuleName, 7, "invalid parameter") + + // ErrOngoingDistribution is returned when a schedule update is attempted during an ongoing distribution. + ErrOngoingDistribution = sdkerrors.Register(ModuleName, 8, "distribution is currently in progress") ) From 5ef1e675f7b457368e0f6c49a67b863df2e44156 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Mon, 2 Mar 2026 10:08:40 +0300 Subject: [PATCH 10/34] remove deprecated score map --- x/pse/keeper/score_map.go | 166 -------------------------------------- 1 file changed, 166 deletions(-) delete mode 100644 x/pse/keeper/score_map.go diff --git a/x/pse/keeper/score_map.go b/x/pse/keeper/score_map.go deleted file mode 100644 index 25dc6664..00000000 --- a/x/pse/keeper/score_map.go +++ /dev/null @@ -1,166 +0,0 @@ -package keeper - -import ( - "context" - - "cosmossdk.io/collections" - addresscodec "cosmossdk.io/core/address" - sdkmath "cosmossdk.io/math" - sdk "github.com/cosmos/cosmos-sdk/types" - - "github.com/tokenize-x/tx-chain/v7/x/pse/types" -) - -type scoreMap struct { - items []struct { - addr sdk.AccAddress - score sdkmath.Int - } - distributionID uint64 - indexMap map[string]int - addressCodec addresscodec.Codec - totalScore sdkmath.Int - excludedAddresses []sdk.AccAddress -} - -func newScoreMap( - distributionID uint64, - addressCodec addresscodec.Codec, - excludedAddressesStr []string, -) (*scoreMap, error) { - excludedAddresses := make([]sdk.AccAddress, len(excludedAddressesStr)) - for i, addr := range excludedAddressesStr { - var err error - excludedAddresses[i], err = addressCodec.StringToBytes(addr) - if err != nil { - return nil, err - } - } - return &scoreMap{ - items: make([]struct { - addr sdk.AccAddress - score sdkmath.Int - }, 0), - distributionID: distributionID, - indexMap: make(map[string]int), - addressCodec: addressCodec, - totalScore: sdkmath.NewInt(0), - excludedAddresses: excludedAddresses, - }, nil -} - -func (m *scoreMap) addScore(addr sdk.AccAddress, value sdkmath.Int) error { - if value.IsZero() { - return nil - } - key, err := m.addressCodec.BytesToString(addr) - if err != nil { - return err - } - idx, found := m.indexMap[key] - if !found { - m.items = append(m.items, struct { - addr sdk.AccAddress - score sdkmath.Int - }{ - addr: addr, - score: value, - }) - m.indexMap[key] = len(m.items) - 1 - } else { - m.items[idx].score = m.items[idx].score.Add(value) - } - - m.totalScore = m.totalScore.Add(value) - return nil -} - -func (m *scoreMap) walk(fn func(addr sdk.AccAddress, score sdkmath.Int) error) error { - for _, pair := range m.items { - if m.isExcludedAddress(pair.addr) { - continue - } - if err := fn(pair.addr, pair.score); err != nil { - return err - } - } - return nil -} - -func (m *scoreMap) iterateAccountScoreSnapshot(ctx context.Context, k Keeper) error { - iter, err := k.AccountScoreSnapshot.Iterate( - ctx, - collections.NewPrefixedPairRange[uint64, sdk.AccAddress](m.distributionID), - ) - if err != nil { - return err - } - defer iter.Close() - for ; iter.Valid(); iter.Next() { - kv, err := iter.KeyValue() - if err != nil { - return err - } - score := kv.Value - delAddr := kv.Key.K2() - if m.isExcludedAddress(delAddr) { - continue - } - err = m.addScore(delAddr, score) - if err != nil { - return err - } - } - - return nil -} - -func (m *scoreMap) iterateDelegationTimeEntries(ctx context.Context, k Keeper) ( - []collections.KeyValue[collections.Triple[uint64, sdk.AccAddress, sdk.ValAddress], types.DelegationTimeEntry], error, -) { - var allDelegationTimeEntries []collections.KeyValue[ - collections.Triple[uint64, sdk.AccAddress, sdk.ValAddress], - types.DelegationTimeEntry, - ] - delegationTimeEntriesIterator, err := k.DelegationTimeEntries.Iterate( - ctx, - collections.NewPrefixedTripleRange[uint64, sdk.AccAddress, sdk.ValAddress](m.distributionID), - ) - if err != nil { - return nil, err - } - defer delegationTimeEntriesIterator.Close() - - for ; delegationTimeEntriesIterator.Valid(); delegationTimeEntriesIterator.Next() { - kv, err := delegationTimeEntriesIterator.KeyValue() - if err != nil { - return nil, err - } - allDelegationTimeEntries = append(allDelegationTimeEntries, kv) - delAddr := kv.Key.K2() - valAddr := kv.Key.K3() - if m.isExcludedAddress(delAddr) { - continue - } - - delegationTimeEntry := kv.Value - delegationScore, err := calculateAddedScore(ctx, k, valAddr, delegationTimeEntry) - if err != nil { - return nil, err - } - err = m.addScore(delAddr, delegationScore) - if err != nil { - return nil, err - } - } - return allDelegationTimeEntries, nil -} - -func (m *scoreMap) isExcludedAddress(addr sdk.AccAddress) bool { - for _, excludedAddress := range m.excludedAddresses { - if excludedAddress.Equals(addr) { - return true - } - } - return false -} From c060559c0510aff9c1d40a16aa4bde107c0d8c78 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Mon, 2 Mar 2026 10:12:13 +0300 Subject: [PATCH 11/34] simplify getCurrentDistributionID to support zero-based IDs --- x/pse/keeper/hooks.go | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/x/pse/keeper/hooks.go b/x/pse/keeper/hooks.go index 1559850d..d5051d4c 100644 --- a/x/pse/keeper/hooks.go +++ b/x/pse/keeper/hooks.go @@ -38,8 +38,8 @@ func (k Keeper) getOngoingDistribution(ctx context.Context) (*types.ScheduledDis // getCurrentDistributionID returns the distribution ID that new entries should be written to. // If an ongoing distribution exists (ID=N is being processed), returns N+1. -// Otherwise returns the next scheduled distribution's ID. -// Returns 0 if no distribution is scheduled and none is ongoing. +// Otherwise returns the next scheduled distribution's ID (zero-value ID when no schedule exists). +// TODO: handle empty distribution schedule — currently returns 0 when no schedule exists func (k Keeper) getCurrentDistributionID(ctx context.Context) (uint64, error) { ongoing, err := k.getOngoingDistribution(ctx) if err != nil { @@ -71,9 +71,6 @@ func (h Hooks) AfterDelegationModified(ctx context.Context, delAddr sdk.AccAddre if err != nil { return err } - if currentID == 0 { - return nil - } isExcluded, err := h.k.IsExcludedAddress(ctx, delAddr) if err != nil { @@ -165,9 +162,6 @@ func (h Hooks) BeforeDelegationRemoved(ctx context.Context, delAddr sdk.AccAddre if err != nil { return err } - if currentID == 0 { - return nil - } isExcluded, err := h.k.IsExcludedAddress(ctx, delAddr) if err != nil { From c309adbaa9d0a870c8e808a9183aa61976e93893 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Mon, 2 Mar 2026 10:13:35 +0300 Subject: [PATCH 12/34] allow zero-based distribution IDs in schedule validation --- x/pse/keeper/params.go | 2 +- x/pse/types/params_test.go | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/x/pse/keeper/params.go b/x/pse/keeper/params.go index 52c1fe2a..86bc4161 100644 --- a/x/pse/keeper/params.go +++ b/x/pse/keeper/params.go @@ -62,7 +62,7 @@ func (k Keeper) UpdateExcludedAddresses( if err != nil { return err } - distributionID := distribution.ID // TODO update to handle distribution ID properly. + distributionID := distribution.ID for _, addrStr := range addressesToRemove { addr, err := k.addressCodec.StringToBytes(addrStr) if err != nil { diff --git a/x/pse/types/params_test.go b/x/pse/types/params_test.go index ef09c3d3..d048e3df 100644 --- a/x/pse/types/params_test.go +++ b/x/pse/types/params_test.go @@ -364,7 +364,7 @@ func TestValidateAllocationSchedule(t *testing.T) { Allocations: createAllModuleAllocations(sdkmath.NewInt(1000)), }, }, - expectErr: false, // TODO: this should be true or removed, bases on the decision on the id zero check. + expectErr: false, // TODO: this should be handled based on the decision on the id zero check. errMsg: "id cannot be zero", }, { From 03acdebd2694446cbe9009c1797f1deb3dad0a28 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Mon, 2 Mar 2026 10:15:12 +0300 Subject: [PATCH 13/34] fix unit tests for zero-based IDs and multi-block phases --- x/pse/keeper/distribute_test.go | 42 ++++++++++++++++++---- x/pse/keeper/distribution_test.go | 10 ++++-- x/pse/keeper/hooks_test.go | 58 ++++++++++++++++++++++++------- 3 files changed, 88 insertions(+), 22 deletions(-) diff --git a/x/pse/keeper/distribute_test.go b/x/pse/keeper/distribute_test.go index 16e6f32e..2c3d5fcf 100644 --- a/x/pse/keeper/distribute_test.go +++ b/x/pse/keeper/distribute_test.go @@ -33,8 +33,8 @@ func TestKeeper_Distribute(t *testing.T) { func(r *runEnv) { distributeAction(r, sdkmath.NewInt(1000)) }, func(r *runEnv) { assertDistributionAction(r, map[*sdk.AccAddress]sdkmath.Int{ - &r.delegators[0]: sdkmath.NewInt(1_100_550), // + 1000 * 1.1 / 3 - &r.delegators[1]: sdkmath.NewInt(900_450), // + 1000 * 0.9 / 3 + &r.delegators[0]: sdkmath.NewInt(1_100_366), // + 1000 * 1.1 / 2 + &r.delegators[1]: sdkmath.NewInt(900_299), // + 1000 * 0.9 / 2 }) }, func(r *runEnv) { assertScoreResetAction(r) }, @@ -193,9 +193,10 @@ func TestKeeper_Distribute(t *testing.T) { ctx, _, err := testApp.BeginNextBlockAtTime(startTime) requireT.NoError(err) runContext := &runEnv{ - testApp: testApp, - ctx: ctx, - requireT: requireT, + testApp: testApp, + ctx: ctx, + requireT: requireT, + currentDistID: tempDistributionID, } // add validators. @@ -349,16 +350,43 @@ func Test_ExcludedAddress_FullLifecycle(t *testing.T) { scheduledDistribution := types.ScheduledDistribution{ ID: distributionID, Timestamp: uint64(ctx.BlockTime().Unix()), + ID: distributionID, + Allocations: []types.ClearingAccountAllocation{{ + ClearingAccount: types.ClearingAccountCommunity, + Amount: amount, + }}, } - balanceBefore := testApp.BankKeeper.GetBalance(ctx, delAddr, bondDenom) - err = pseKeeper.DistributeCommunityPSE(ctx, bondDenom, amount, scheduledDistribution) + err = pseKeeper.OngoingDistribution.Set(ctx, scheduledDistribution) requireT.NoError(err) + balanceBefore := testApp.BankKeeper.GetBalance(ctx, delAddr, bondDenom) + for { + done, err := pseKeeper.ProcessPhase1ScoreConversion(ctx, scheduledDistribution) + requireT.NoError(err) + if done { + break + } + } + for { + done, err := pseKeeper.ProcessPhase2TokenDistribution(ctx, scheduledDistribution, bondDenom) + requireT.NoError(err) + if done { + break + } + } balanceAfter := testApp.BankKeeper.GetBalance(ctx, delAddr, bondDenom) requireT.Equal( balanceBefore.Amount.String(), balanceAfter.Amount.String(), "Excluded address should receive no rewards", ) + // After distribution, entries migrated from distributionID to distributionID+1. + // Save a new schedule so hooks and UpdateExcludedAddresses can find it. + distributionID++ + err = pseKeeper.SaveDistributionSchedule(ctx, []types.ScheduledDistribution{ + {Timestamp: distributionID, ID: distributionID}, + }) + requireT.NoError(err) + // Step 6: Verify excluded delegator can fully undelegate after distribution msgUndel := &stakingtypes.MsgUndelegate{ DelegatorAddress: delAddr.String(), diff --git a/x/pse/keeper/distribution_test.go b/x/pse/keeper/distribution_test.go index fa386454..65c8cef6 100644 --- a/x/pse/keeper/distribution_test.go +++ b/x/pse/keeper/distribution_test.go @@ -97,8 +97,14 @@ func TestDistribution_GenesisRebuild(t *testing.T) { // Process first distribution ctx = ctx.WithBlockTime(time.Unix(int64(time1)+10, 0)) ctx = ctx.WithBlockHeight(100) - err = pseKeeper.ProcessNextDistribution(ctx) - requireT.NoError(err) + for range 10 { + err = pseKeeper.ProcessNextDistribution(ctx) + requireT.NoError(err) + _, oErr := pseKeeper.OngoingDistribution.Get(ctx) + if oErr != nil { + break + } + } // Export genesis genesisState, err := pseKeeper.ExportGenesis(ctx) diff --git a/x/pse/keeper/hooks_test.go b/x/pse/keeper/hooks_test.go index b5171b02..91deb7f7 100644 --- a/x/pse/keeper/hooks_test.go +++ b/x/pse/keeper/hooks_test.go @@ -199,9 +199,10 @@ func TestKeeper_Hooks(t *testing.T) { testApp := simapp.New() ctx := testApp.NewContext(false) runContext := &runEnv{ - testApp: testApp, - ctx: ctx, - requireT: requireT, + testApp: testApp, + ctx: ctx, + requireT: requireT, + currentDistID: tempDistributionID, } err := testApp.PSEKeeper.SaveDistributionSchedule(ctx, []types.ScheduledDistribution{ @@ -247,16 +248,17 @@ func TestKeeper_Hooks(t *testing.T) { } type runEnv struct { - testApp *simapp.App - ctx sdk.Context - delegators []sdk.AccAddress - validators []sdk.ValAddress - requireT *require.Assertions + testApp *simapp.App + ctx sdk.Context + delegators []sdk.AccAddress + validators []sdk.ValAddress + requireT *require.Assertions + currentDistID uint64 } func assertScoreAction(r *runEnv, delAddr sdk.AccAddress, expectedScore sdkmath.Int) { score, err := r.testApp.PSEKeeper.GetDelegatorScore( - r.ctx, tempDistributionID, delAddr, + r.ctx, r.currentDistID, delAddr, ) r.requireT.NoError(err) r.requireT.Equal(expectedScore, score) @@ -285,7 +287,9 @@ func assertCommunityPoolBalanceAction(r *runEnv, expectedBalance sdkmath.Int) { } func assertScoreResetAction(r *runEnv) { - scoreRanger := collections.NewPrefixedPairRange[uint64, sdk.AccAddress](tempDistributionID) + // After cleanup, score snapshots at the previous distribution ID should be cleared. + prevID := r.currentDistID - 1 + scoreRanger := collections.NewPrefixedPairRange[uint64, sdk.AccAddress](prevID) err := r.testApp.PSEKeeper.AccountScoreSnapshot.Walk(r.ctx, scoreRanger, func(key collections.Pair[uint64, sdk.AccAddress], value sdkmath.Int) (bool, error) { r.requireT.Equal(sdkmath.NewInt(0), value) @@ -293,8 +297,9 @@ func assertScoreResetAction(r *runEnv) { }) r.requireT.NoError(err) + // Entries should exist at the current distribution ID (migrated during Phase 1). blockTimeUnixSeconds := r.ctx.BlockTime().Unix() - entriesRanger := collections.NewPrefixedTripleRange[uint64, sdk.AccAddress, sdk.ValAddress](tempDistributionID) + entriesRanger := collections.NewPrefixedTripleRange[uint64, sdk.AccAddress, sdk.ValAddress](r.currentDistID) err = r.testApp.PSEKeeper.DelegationTimeEntries.Walk(r.ctx, entriesRanger, func( key collections.Triple[uint64, sdk.AccAddress, sdk.ValAddress], value types.DelegationTimeEntry, @@ -372,10 +377,37 @@ func distributeAction(r *runEnv, amount sdkmath.Int) { r.requireT.NoError(err) scheduledDistribution := types.ScheduledDistribution{ Timestamp: uint64(r.ctx.BlockTime().Unix()), - ID: tempDistributionID, + ID: r.currentDistID, + Allocations: []types.ClearingAccountAllocation{{ + ClearingAccount: types.ClearingAccountCommunity, + Amount: amount, + }}, } - err = r.testApp.PSEKeeper.DistributeCommunityPSE(r.ctx, bondDenom, amount, scheduledDistribution) + + // Set OngoingDistribution to simulate EndBlocker starting multi-block processing. + err = r.testApp.PSEKeeper.OngoingDistribution.Set(r.ctx, scheduledDistribution) r.requireT.NoError(err) + + // Run Phase 1 until done. + for { + done, err := r.testApp.PSEKeeper.ProcessPhase1ScoreConversion(r.ctx, scheduledDistribution) + r.requireT.NoError(err) + if done { + break + } + } + + // Run Phase 2 until done. + for { + done, err := r.testApp.PSEKeeper.ProcessPhase2TokenDistribution(r.ctx, scheduledDistribution, bondDenom) + r.requireT.NoError(err) + if done { + break + } + } + + // Advance to next distribution ID (Phase 1 migrated entries to currentDistID+1). + r.currentDistID++ } func mintAndSendCoin(r *runEnv, recipient sdk.AccAddress, coins sdk.Coins) { From 4f1b21bf5d438a49c76c3f2d7421882963e85185 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Mon, 2 Mar 2026 10:16:00 +0300 Subject: [PATCH 14/34] remove deprecated single-block distribute functions --- x/pse/keeper/distribute.go | 109 ------------------------------ x/pse/keeper/distribution.go | 124 ----------------------------------- 2 files changed, 233 deletions(-) diff --git a/x/pse/keeper/distribute.go b/x/pse/keeper/distribute.go index 33e0cdff..26999cf7 100644 --- a/x/pse/keeper/distribute.go +++ b/x/pse/keeper/distribute.go @@ -287,115 +287,6 @@ func (k Keeper) addToDistributedAmount(ctx context.Context, distributionID uint6 return k.DistributedAmount.Set(ctx, distributionID, current.Add(amount)) } -// DistributeCommunityPSE distributes the total community PSE amount to all delegators based on their score. -func (k Keeper) DistributeCommunityPSE( - ctx context.Context, - bondDenom string, - totalPSEAmount sdkmath.Int, - scheduledDistribution types.ScheduledDistribution, -) error { - scheduledAt := scheduledDistribution.Timestamp - // TODO update to use distribution ID, also consider period splits - distributionID := scheduledDistribution.ID - // iterate all delegation time entries and calculate uncalculated score. - params, err := k.GetParams(ctx) - if err != nil { - return err - } - finalScoreMap, err := newScoreMap(distributionID, k.addressCodec, params.ExcludedAddresses) - if err != nil { - return err - } - - allDelegationTimeEntries, err := finalScoreMap.iterateDelegationTimeEntries(ctx, k) - if err != nil { - return err - } - - // add uncalculated score to account score snapshot and total score per delegator. - // it calculates the score from the last delegation time entry up to the current block time, which - // is not included in the score snapshot calculations. - err = finalScoreMap.iterateAccountScoreSnapshot(ctx, k) - if err != nil { - return err - } - - // Clear all account score snapshots. - // Excluded addresses should not have snapshots (cleared when added to exclusion list), - // but we clear unconditionally for all addresses. - // TODO review all the logic for score reset - if err := k.AccountScoreSnapshot.Clear( - ctx, - collections.NewPrefixedPairRange[uint64, sdk.AccAddress](distributionID), - ); err != nil { - return err - } - - // reset all delegation time entries LastChangedUnixSec to the current block time. - err = k.DelegationTimeEntries.Clear( - ctx, - collections.NewPrefixedTripleRange[uint64, sdk.AccAddress, sdk.ValAddress](distributionID), - ) - if err != nil { - return err - } - currentBlockTime := sdk.UnwrapSDKContext(ctx).BlockTime().Unix() - for _, kv := range allDelegationTimeEntries { - kv.Value.LastChangedUnixSec = currentBlockTime - // TODO review all the logic for score reset - key := collections.Join3(distributionID+1, kv.Key.K2(), kv.Key.K3()) - err = k.DelegationTimeEntries.Set(ctx, key, kv.Value) - if err != nil { - return err - } - } - - // distribute total pse coin based on per delegator score. - totalPSEScore := finalScoreMap.totalScore - - // leftover is the amount of pse coin that is not distributed to any delegator. - // It will be sent to CommunityPool. - // there are 2 sources of leftover: - // 1. rounding errors due to division. - // 2. some delegators have no delegation. - leftover := totalPSEAmount - sdkCtx := sdk.UnwrapSDKContext(ctx) - if totalPSEScore.IsPositive() { - err = finalScoreMap.walk(func(addr sdk.AccAddress, score sdkmath.Int) error { - userAmount := totalPSEAmount.Mul(score).Quo(totalPSEScore) - distributedAmount, err := k.distributeToDelegator(ctx, addr, userAmount, bondDenom) - if err != nil { - return err - } - leftover = leftover.Sub(distributedAmount) - if err := sdkCtx.EventManager().EmitTypedEvent(&types.EventCommunityDistributed{ - DelegatorAddress: addr.String(), - Score: score, - TotalPseScore: totalPSEScore, - Amount: userAmount, - ScheduledAt: scheduledAt, - }); err != nil { - sdkCtx.Logger().Error("failed to emit community distributed event", "error", err) - } - return nil - }) - if err != nil { - return err - } - } - - // send leftover to CommunityPool. - if leftover.IsPositive() { - pseModuleAddress := k.accountKeeper.GetModuleAddress(types.ClearingAccountCommunity) - err = k.distributionKeeper.FundCommunityPool(ctx, sdk.NewCoins(sdk.NewCoin(bondDenom, leftover)), pseModuleAddress) - if err != nil { - return err - } - } - - return nil -} - func (k Keeper) distributeToDelegator( ctx context.Context, delAddr sdk.AccAddress, amount sdkmath.Int, bondDenom string, ) (sdkmath.Int, error) { diff --git a/x/pse/keeper/distribution.go b/x/pse/keeper/distribution.go index 8c0224a1..a3b6c1b3 100644 --- a/x/pse/keeper/distribution.go +++ b/x/pse/keeper/distribution.go @@ -238,130 +238,6 @@ func (k Keeper) PeekNextAllocationSchedule(ctx context.Context) (types.Scheduled return scheduledDist, shouldProcess, nil } -// distributeAllocatedTokens transfers tokens from clearing accounts to their mapped recipients. -// Processes all allocations within a single scheduled distribution. -// Any transfer failure indicates a state invariant violation (insufficient balance or invalid recipient). -func (k Keeper) distributeAllocatedTokens( - ctx context.Context, - timestamp uint64, - bondDenom string, - clearingAccountMappings []types.ClearingAccountMapping, - scheduledDistribution types.ScheduledDistribution, -) error { - sdkCtx := sdk.UnwrapSDKContext(ctx) - // Transfer tokens for each allocation in this distribution period - for _, allocation := range scheduledDistribution.Allocations { - if allocation.Amount.IsZero() { - continue - } - - // Community clearing account has different distribution logic - if allocation.ClearingAccount == types.ClearingAccountCommunity { - if err := k.DistributeCommunityPSE(ctx, bondDenom, allocation.Amount, scheduledDistribution); err != nil { - return errorsmod.Wrapf( - types.ErrTransferFailed, - "failed to distribute Community clearing account allocation: %v", - err, - ) - } - continue - } - - // Find the recipient addresses mapped to this clearing account - // Note: Community clearing account is handled above and doesn't need a mapping. - // Mappings are validated on update and genesis, so they are guaranteed to exist. - var recipientAddrs []string - for _, mapping := range clearingAccountMappings { - if mapping.ClearingAccount == allocation.ClearingAccount { - recipientAddrs = mapping.RecipientAddresses - break - } - } - - // Distribution Precision Handling: - // The allocation amount is split equally among all recipients using integer division. - // Any remainder from division is sent to the community pool to ensure: - // - Each recipient receives exactly: allocation.Amount / numRecipients (base amount) - // - Remainder (if any) goes to community pool for ecosystem benefit - // This guarantees fair distribution and no tokens are lost - numRecipients := sdkmath.NewInt(int64(len(recipientAddrs))) - if numRecipients.IsZero() { - return errorsmod.Wrapf( - types.ErrTransferFailed, - "no recipients found for clearing account '%s'", - allocation.ClearingAccount, - ) - } - amountPerRecipient := allocation.Amount.Quo(numRecipients) - remainder := allocation.Amount.Mod(numRecipients) - - // Transfer tokens to each recipient - for _, recipientAddr := range recipientAddrs { - // Convert recipient address string to SDK account address - // Safe to use Must* because addresses are validated at genesis/update time - recipient := sdk.MustAccAddressFromBech32(recipientAddr) - - // Each recipient gets equal base amount - coinsToSend := sdk.NewCoins(sdk.NewCoin(bondDenom, amountPerRecipient)) - - // Transfer tokens from clearing account to recipient - if err := k.bankKeeper.SendCoinsFromModuleToAccount( - ctx, - allocation.ClearingAccount, - recipient, - coinsToSend, - ); err != nil { - return errorsmod.Wrapf( - types.ErrTransferFailed, - "failed to transfer from clearing account '%s' to recipient '%s': %v", - allocation.ClearingAccount, - recipientAddr, - err, - ) - } - } - - // Send any remainder to community pool - if !remainder.IsZero() { - clearingAccountAddr := k.accountKeeper.GetModuleAddress(allocation.ClearingAccount) - remainderCoins := sdk.NewCoins(sdk.NewCoin(bondDenom, remainder)) - if err := k.distributionKeeper.FundCommunityPool(ctx, remainderCoins, clearingAccountAddr); err != nil { - return errorsmod.Wrapf( - types.ErrTransferFailed, - "failed to send remainder to community pool from clearing account '%s': %v", - allocation.ClearingAccount, - err, - ) - } - - sdkCtx.Logger().Info("sent distribution remainder to community pool", - "clearing_account", allocation.ClearingAccount, - "remainder", remainder.String()) - } - - // Emit single allocation completed event with recipient list, per-recipient amount, and community pool amount - if err := sdkCtx.EventManager().EmitTypedEvent(&types.EventAllocationDistributed{ - ClearingAccount: allocation.ClearingAccount, - RecipientAddresses: recipientAddrs, - AmountPerRecipient: amountPerRecipient, - CommunityPoolAmount: remainder, - ScheduledAt: timestamp, - TotalAmount: allocation.Amount, - }); err != nil { - sdkCtx.Logger().Error("failed to emit allocation completed event", "error", err) - } - - sdkCtx.Logger().Info("allocated tokens", - "clearing_account", allocation.ClearingAccount, - "recipients", recipientAddrs, - "total_amount", allocation.Amount.String(), - "amount_per_recipient", amountPerRecipient.String(), - "community_pool_amount", remainder.String()) - } - - return nil -} - // SaveDistributionSchedule persists the distribution schedule to blockchain state. // Each scheduled distribution is stored in the AllocationSchedule map, indexed by its ID. func (k Keeper) SaveDistributionSchedule(ctx context.Context, schedule []types.ScheduledDistribution) error { From dd72208e631440c7e82f4aa64a0eb3b1a7618e63 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Mon, 2 Mar 2026 11:05:56 +0300 Subject: [PATCH 15/34] fix lint --- x/pse/keeper/delegation.go | 8 +- x/pse/keeper/distribute.go | 4 +- x/pse/keeper/distribution.go | 15 ++-- x/pse/keeper/hooks.go | 160 ++++++++++++++++----------------- x/pse/keeper/keeper.go | 8 +- x/pse/types/distribution.pb.go | 9 +- 6 files changed, 102 insertions(+), 102 deletions(-) diff --git a/x/pse/keeper/delegation.go b/x/pse/keeper/delegation.go index 141f22f1..7350fb3e 100644 --- a/x/pse/keeper/delegation.go +++ b/x/pse/keeper/delegation.go @@ -56,7 +56,9 @@ func (k Keeper) GetDelegatorScore( } // SetDelegatorScore sets the score for a delegator. -func (k Keeper) SetDelegatorScore(ctx context.Context, distributionID uint64, delAddr sdk.AccAddress, score sdkmath.Int) error { +func (k Keeper) SetDelegatorScore( + ctx context.Context, distributionID uint64, delAddr sdk.AccAddress, score sdkmath.Int, +) error { key := collections.Join(distributionID, delAddr) return k.AccountScoreSnapshot.Set(ctx, key, score) } @@ -68,7 +70,9 @@ func (k Keeper) RemoveDelegatorScore(ctx context.Context, distributionID uint64, } // addToScore atomically adds a score value to a delegator's score snapshot. -func (k Keeper) addToScore(ctx context.Context, distributionID uint64, delAddr sdk.AccAddress, score sdkmath.Int) error { +func (k Keeper) addToScore( + ctx context.Context, distributionID uint64, delAddr sdk.AccAddress, score sdkmath.Int, +) error { if score.IsZero() { return nil } diff --git a/x/pse/keeper/distribute.go b/x/pse/keeper/distribute.go index 26999cf7..413af6ed 100644 --- a/x/pse/keeper/distribute.go +++ b/x/pse/keeper/distribute.go @@ -134,7 +134,9 @@ func (k Keeper) computeTotalScore(ctx context.Context, distributionID uint64) er // // When all delegators have been processed, sends leftover (rounding errors + undelegated users) to the community pool. // Returns true when distribution is complete and all state has been cleaned up. -func (k Keeper) ProcessPhase2TokenDistribution(ctx context.Context, ongoing types.ScheduledDistribution, bondDenom string) (bool, error) { +func (k Keeper) ProcessPhase2TokenDistribution( + ctx context.Context, ongoing types.ScheduledDistribution, bondDenom string, +) (bool, error) { prevID := ongoing.ID totalPSEAmount := getCommunityAllocationAmount(ongoing) diff --git a/x/pse/keeper/distribution.go b/x/pse/keeper/distribution.go index a3b6c1b3..e6c59fa7 100644 --- a/x/pse/keeper/distribution.go +++ b/x/pse/keeper/distribution.go @@ -21,12 +21,12 @@ import ( // c. Else, no community allocation, non-community distribution is already done, remove from AllocationSchedule func (k Keeper) ProcessNextDistribution(ctx context.Context) error { // Resume ongoing multi-block distribution if one is in progress. - ongoing, err := k.getOngoingDistribution(ctx) + ongoing, found, err := k.getOngoingDistribution(ctx) if err != nil { return err } - if ongoing != nil { - return k.resumeOngoingDistribution(ctx, *ongoing) + if found { + return k.resumeOngoingDistribution(ctx, ongoing) } // No ongoing distribution — check if next scheduled distribution is due. @@ -289,12 +289,15 @@ func (k Keeper) UpdateDistributionSchedule( } // Reject if a multi-block distribution is in progress. - ongoing, err := k.getOngoingDistribution(ctx) + ongoing, ongoingFound, err := k.getOngoingDistribution(ctx) if err != nil { return err } - if ongoing != nil { - return errorsmod.Wrapf(types.ErrOngoingDistribution, "cannot update schedule while distribution %d is in progress", ongoing.ID) + if ongoingFound { + return errorsmod.Wrapf( + types.ErrOngoingDistribution, + "cannot update schedule while distribution %d is in progress", ongoing.ID, + ) } // Validate minimum gap between distributions diff --git a/x/pse/keeper/hooks.go b/x/pse/keeper/hooks.go index d5051d4c..369b5922 100644 --- a/x/pse/keeper/hooks.go +++ b/x/pse/keeper/hooks.go @@ -24,28 +24,28 @@ func (k Keeper) Hooks() Hooks { return Hooks{k} } -// getOngoingDistribution returns the ongoing distribution if one exists, or nil if not. -func (k Keeper) getOngoingDistribution(ctx context.Context) (*types.ScheduledDistribution, error) { +// getOngoingDistribution returns the ongoing distribution if one exists. +func (k Keeper) getOngoingDistribution(ctx context.Context) (types.ScheduledDistribution, bool, error) { ongoing, err := k.OngoingDistribution.Get(ctx) if errors.Is(err, collections.ErrNotFound) { - return nil, nil + return types.ScheduledDistribution{}, false, nil } if err != nil { - return nil, err + return types.ScheduledDistribution{}, false, err } - return &ongoing, nil + return ongoing, true, nil } // getCurrentDistributionID returns the distribution ID that new entries should be written to. // If an ongoing distribution exists (ID=N is being processed), returns N+1. // Otherwise returns the next scheduled distribution's ID (zero-value ID when no schedule exists). -// TODO: handle empty distribution schedule — currently returns 0 when no schedule exists +// TODO: handle empty distribution schedule — currently returns 0 when no schedule exists. func (k Keeper) getCurrentDistributionID(ctx context.Context) (uint64, error) { - ongoing, err := k.getOngoingDistribution(ctx) + ongoing, found, err := k.getOngoingDistribution(ctx) if err != nil { return 0, err } - if ongoing != nil { + if found { return ongoing.ID + 1, nil } @@ -58,7 +58,7 @@ func (k Keeper) getCurrentDistributionID(ctx context.Context) (uint64, error) { // AfterDelegationModified implements the staking hooks interface. // Handles 3 scenarios based on where the delegator's entry exists: -// - Scenario 1: Entry in prevID (ongoing distribution in progress) — +// - Scenario 1: Entry in prevID (ongoing distribution in progress). // - Scenario 2: Entry in currentID — normal score calculation. // - Scenario 3: No entry — create new entry, no score. func (h Hooks) AfterDelegationModified(ctx context.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress) error { @@ -84,53 +84,24 @@ func (h Hooks) AfterDelegationModified(ctx context.Context, delAddr sdk.AccAddre // Scenario 1: Entry exists in previous distribution (ongoing distribution in progress). // Split score at distribution timestamp, move entry to currentID. - ongoing, err := h.k.getOngoingDistribution(ctx) + ongoing, ongoingFound, err := h.k.getOngoingDistribution(ctx) if err != nil { return err } - if ongoing != nil { - prevID := ongoing.ID - prevEntry, err := h.k.GetDelegationTimeEntry(ctx, prevID, valAddr, delAddr) - if err == nil { - distTimestamp := int64(ongoing.Timestamp) - - // Score for previous period: lastChanged -> distribution timestamp - prevScore, err := calculateScoreAtTimestamp(ctx, h.k, valAddr, prevEntry, distTimestamp) - if err != nil { - return err - } - if err := h.k.addToScore(ctx, prevID, delAddr, prevScore); err != nil { - return err - } - - // Score for current period: distribution timestamp -> now (old shares still active) - currentPeriodEntry := types.DelegationTimeEntry{ - LastChangedUnixSec: distTimestamp, - Shares: prevEntry.Shares, - } - currentScore, err := calculateScoreAtTimestamp(ctx, h.k, valAddr, currentPeriodEntry, blockTime) - if err != nil { - return err - } - if err := h.k.addToScore(ctx, currentID, delAddr, currentScore); err != nil { - return err - } - - // Delete from prevID, create in currentID with new shares - if err := h.k.RemoveDelegationTimeEntry(ctx, prevID, valAddr, delAddr); err != nil { - return err - } + if ongoingFound { + handled, err := h.migrateOngoingEntry(ctx, ongoing, currentID, delAddr, valAddr, blockTime) + if err != nil { + return err + } + if handled { return h.k.SetDelegationTimeEntry(ctx, currentID, valAddr, delAddr, types.DelegationTimeEntry{ LastChangedUnixSec: blockTime, Shares: delegation.Shares, }) } - if !errors.Is(err, collections.ErrNotFound) { - return err - } } - // Scenario 2: Entry exists in current distribution + // Scenario 2: Entry exists in current distribution. currentEntry, err := h.k.GetDelegationTimeEntry(ctx, currentID, valAddr, delAddr) if err == nil { score, err := calculateAddedScore(ctx, h.k, valAddr, currentEntry) @@ -149,7 +120,7 @@ func (h Hooks) AfterDelegationModified(ctx context.Context, delAddr sdk.AccAddre return err } - // Scenario 3: No entry - create new in currentID (no score, duration = 0) + // Scenario 3: No entry - create new in currentID (no score, duration = 0). return h.k.SetDelegationTimeEntry(ctx, currentID, valAddr, delAddr, types.DelegationTimeEntry{ LastChangedUnixSec: blockTime, Shares: delegation.Shares, @@ -173,48 +144,22 @@ func (h Hooks) BeforeDelegationRemoved(ctx context.Context, delAddr sdk.AccAddre blockTime := sdk.UnwrapSDKContext(ctx).BlockTime().Unix() - // Scenario 1: Entry exists in previous distribution (ongoing) - ongoing, err := h.k.getOngoingDistribution(ctx) + // Scenario 1: Entry exists in previous distribution (ongoing). + ongoing, ongoingFound, err := h.k.getOngoingDistribution(ctx) if err != nil { return err } - if ongoing != nil { - prevID := ongoing.ID - prevEntry, err := h.k.GetDelegationTimeEntry(ctx, prevID, valAddr, delAddr) - if err == nil { - distTimestamp := int64(ongoing.Timestamp) - - // Score for previous period: lastChanged -> distribution timestamp - prevScore, err := calculateScoreAtTimestamp(ctx, h.k, valAddr, prevEntry, distTimestamp) - if err != nil { - return err - } - if err := h.k.addToScore(ctx, prevID, delAddr, prevScore); err != nil { - return err - } - - // Score for current period: distribution timestamp -> now - currentPeriodEntry := types.DelegationTimeEntry{ - LastChangedUnixSec: distTimestamp, - Shares: prevEntry.Shares, - } - currentScore, err := calculateScoreAtTimestamp(ctx, h.k, valAddr, currentPeriodEntry, blockTime) - if err != nil { - return err - } - if err := h.k.addToScore(ctx, currentID, delAddr, currentScore); err != nil { - return err - } - - // Delete from prevID (delegation removed) - return h.k.RemoveDelegationTimeEntry(ctx, prevID, valAddr, delAddr) - } - if !errors.Is(err, collections.ErrNotFound) { + if ongoingFound { + handled, err := h.migrateOngoingEntry(ctx, ongoing, currentID, delAddr, valAddr, blockTime) + if err != nil { return err } + if handled { + return h.k.RemoveDelegationTimeEntry(ctx, ongoing.ID, valAddr, delAddr) + } } - // Scenario 2: Entry exists in current distribution + // Scenario 2: Entry exists in current distribution. currentEntry, err := h.k.GetDelegationTimeEntry(ctx, currentID, valAddr, delAddr) if err == nil { score, err := calculateAddedScore(ctx, h.k, valAddr, currentEntry) @@ -230,12 +175,12 @@ func (h Hooks) BeforeDelegationRemoved(ctx context.Context, delAddr sdk.AccAddre return err } - // Scenario 3: No entry + // Scenario 3: No entry. return nil } // calculateScoreAtTimestamp calculates the score for a delegation entry up to a specific timestamp. -// score = tokens × (atTimestamp - lastChanged) +// score = tokens × (atTimestamp - lastChanged). func calculateScoreAtTimestamp( ctx context.Context, keeper Keeper, @@ -314,3 +259,50 @@ func (h Hooks) AfterValidatorBeginUnbonding(_ context.Context, _ sdk.ConsAddress func (h Hooks) AfterUnbondingInitiated(_ context.Context, _ uint64) error { return nil } + +// migrateOngoingEntry handles a delegation entry that still lives in the previous (ongoing) distribution. +// It calculates score for both the prev and current periods, removes the entry from prevID, +// and returns true if the entry was found and processed. +func (h Hooks) migrateOngoingEntry( + ctx context.Context, + ongoing types.ScheduledDistribution, + currentID uint64, + delAddr sdk.AccAddress, + valAddr sdk.ValAddress, + blockTime int64, +) (bool, error) { + prevID := ongoing.ID + prevEntry, err := h.k.GetDelegationTimeEntry(ctx, prevID, valAddr, delAddr) + if errors.Is(err, collections.ErrNotFound) { + return false, nil + } + if err != nil { + return false, err + } + + distTimestamp := int64(ongoing.Timestamp) + + // Score for previous period: lastChanged -> distribution timestamp. + prevScore, err := calculateScoreAtTimestamp(ctx, h.k, valAddr, prevEntry, distTimestamp) + if err != nil { + return false, err + } + if err := h.k.addToScore(ctx, prevID, delAddr, prevScore); err != nil { + return false, err + } + + // Score for current period: distribution timestamp -> now. + currentPeriodEntry := types.DelegationTimeEntry{ + LastChangedUnixSec: distTimestamp, + Shares: prevEntry.Shares, + } + currentScore, err := calculateScoreAtTimestamp(ctx, h.k, valAddr, currentPeriodEntry, blockTime) + if err != nil { + return false, err + } + if err := h.k.addToScore(ctx, currentID, delAddr, currentScore); err != nil { + return false, err + } + + return true, nil +} diff --git a/x/pse/keeper/keeper.go b/x/pse/keeper/keeper.go index 4d07f0ef..cc8be6bc 100644 --- a/x/pse/keeper/keeper.go +++ b/x/pse/keeper/keeper.go @@ -37,10 +37,10 @@ type Keeper struct { types.DelegationTimeEntry, ] AccountScoreSnapshot collections.Map[collections.Pair[uint64, sdk.AccAddress], sdkmath.Int] - AllocationSchedule collections.Map[uint64, types.ScheduledDistribution] // Map: ID -> ScheduledDistribution - TotalScore collections.Map[uint64, sdkmath.Int] // Map: ID -> total accumulated score - OngoingDistribution collections.Item[types.ScheduledDistribution] // Currently processing distribution - DistributedAmount collections.Map[uint64, sdkmath.Int] // Map: ID -> cumulative distributed amount + AllocationSchedule collections.Map[uint64, types.ScheduledDistribution] // Map: ID -> ScheduledDistribution + TotalScore collections.Map[uint64, sdkmath.Int] // Map: ID -> total accumulated score + OngoingDistribution collections.Item[types.ScheduledDistribution] // Currently processing distribution + DistributedAmount collections.Map[uint64, sdkmath.Int] // Map: ID -> cumulative distributed amount DistributionDisabled collections.Item[bool] } diff --git a/x/pse/types/distribution.pb.go b/x/pse/types/distribution.pb.go index f1055c0a..ef56433c 100644 --- a/x/pse/types/distribution.pb.go +++ b/x/pse/types/distribution.pb.go @@ -4,15 +4,14 @@ package types import ( - fmt "fmt" - io "io" - math "math" - math_bits "math/bits" - cosmossdk_io_math "cosmossdk.io/math" + fmt "fmt" _ "github.com/cosmos/cosmos-proto" _ "github.com/cosmos/gogoproto/gogoproto" proto "github.com/cosmos/gogoproto/proto" + io "io" + math "math" + math_bits "math/bits" ) // Reference imports to suppress errors if they are not otherwise used. From ba9347002a9253bddbde956de034075db8d0c85f Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Mon, 2 Mar 2026 11:30:11 +0300 Subject: [PATCH 16/34] fix integration test --- integration-tests/modules/pse_test.go | 30 ++++++++++++++++++--------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/integration-tests/modules/pse_test.go b/integration-tests/modules/pse_test.go index dad5b19f..b0e8b57b 100644 --- a/integration-tests/modules/pse_test.go +++ b/integration-tests/modules/pse_test.go @@ -251,9 +251,7 @@ func TestPSEDistribution(t *testing.T) { requireT.NoError(err) t.Logf("Distribution 1 at height: %d", height) - scheduledDistributions, err := getScheduledDistribution(ctx, chain) - requireT.NoError(err) - requireT.Len(scheduledDistributions, 2) + awaitScheduleCount(ctx, t, chain, 2) balancesBefore, scoresBefore, totalScore := getAllDelegatorInfo(ctx, t, chain, height-1) balancesAfter, _, _ := getAllDelegatorInfo(ctx, t, chain, height) @@ -298,9 +296,7 @@ func TestPSEDistribution(t *testing.T) { requireT.NoError(err) t.Logf("Distribution 2 at height: %d", height) - scheduledDistributions, err = getScheduledDistribution(ctx, chain) - requireT.NoError(err) - requireT.Len(scheduledDistributions, 1) + awaitScheduleCount(ctx, t, chain, 1) balancesBefore, scoresBefore, totalScore = getAllDelegatorInfo(ctx, t, chain, height-1) balancesAfter, _, _ = getAllDelegatorInfo(ctx, t, chain, height) @@ -336,9 +332,7 @@ func TestPSEDistribution(t *testing.T) { requireT.NoError(err) t.Logf("Distribution 3 at height: %d", height) - scheduledDistributions, err = getScheduledDistribution(ctx, chain) - requireT.NoError(err) - requireT.Empty(scheduledDistributions) + awaitScheduleCount(ctx, t, chain, 0) balancesBefore, scoresBefore, totalScore = getAllDelegatorInfo(ctx, t, chain, height-1) balancesAfter, _, _ = getAllDelegatorInfo(ctx, t, chain, height) @@ -836,7 +830,7 @@ func awaitScheduledDistributionEvent( ) (int64, communityDistributedEvent, error) { var observedHeight int64 err := chain.AwaitState(ctx, func(ctx context.Context) error { - query := fmt.Sprintf("tx.pse.v1.EventAllocationDistributed.mode='EndBlock' AND block.height>%d", startHeight) + query := fmt.Sprintf("tx.pse.v1.EventCommunityDistributed.mode='EndBlock' AND block.height>%d", startHeight) blocks, err := chain.ClientContext.RPCClient().BlockSearch(ctx, query, nil, nil, "") if err != nil { return err @@ -880,6 +874,22 @@ func getScheduledDistribution( return pseResponse.ScheduledDistributions, nil } +func awaitScheduleCount(ctx context.Context, t *testing.T, chain integration.TXChain, expectedCount int) { + t.Helper() + requireT := require.New(t) + err := chain.AwaitState(ctx, func(ctx context.Context) error { + dist, err := getScheduledDistribution(ctx, chain) + if err != nil { + return err + } + if len(dist) != expectedCount { + return fmt.Errorf("expected %d scheduled distributions, got %d", expectedCount, len(dist)) + } + return nil + }, integration.WithAwaitStateTimeout(10*time.Second)) + requireT.NoError(err) +} + func removeAttributeFromEvent(events []tmtypes.Event, key string) []tmtypes.Event { newEvents := make([]tmtypes.Event, 0, len(events)) for _, event := range events { From 6224afd0d9b006c1eef526a4ef19932e65762b7a Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Mon, 2 Mar 2026 13:10:20 +0300 Subject: [PATCH 17/34] add v1->v2 state migration for PSE collection key changes --- x/pse/keeper/migrations.go | 21 +++- x/pse/migrations/v2/migrate.go | 191 +++++++++++++++++++++++++++++++++ x/pse/module.go | 7 +- 3 files changed, 215 insertions(+), 4 deletions(-) create mode 100644 x/pse/migrations/v2/migrate.go diff --git a/x/pse/keeper/migrations.go b/x/pse/keeper/migrations.go index 47974b8e..fde99103 100644 --- a/x/pse/keeper/migrations.go +++ b/x/pse/keeper/migrations.go @@ -1,10 +1,25 @@ package keeper -// Migrator is a struct for handling in-place store migrations. +import ( + sdk "github.com/cosmos/cosmos-sdk/types" + + v2 "github.com/tokenize-x/tx-chain/v7/x/pse/migrations/v2" +) + +// Migrator handles in-place store migrations for the PSE module. type Migrator struct { + keeper Keeper } // NewMigrator returns a new Migrator. -func NewMigrator() Migrator { - return Migrator{} +func NewMigrator(keeper Keeper) Migrator { + return Migrator{keeper: keeper} +} + +// Migrate1to2 migrates the store from v1 to v2. +// Key changes: +// - DelegationTimeEntries: Pair[AccAddress, ValAddress] -> Triple[uint64, AccAddress, ValAddress]. +// - AccountScoreSnapshot: AccAddress -> Pair[uint64, AccAddress]. +func (m Migrator) Migrate1to2(ctx sdk.Context) error { + return v2.MigrateStore(ctx, m.keeper.storeService, m.keeper.cdc) } diff --git a/x/pse/migrations/v2/migrate.go b/x/pse/migrations/v2/migrate.go new file mode 100644 index 00000000..e4548e5d --- /dev/null +++ b/x/pse/migrations/v2/migrate.go @@ -0,0 +1,191 @@ +package v2 + +import ( + "context" + + "cosmossdk.io/collections" + sdkstore "cosmossdk.io/core/store" + sdkmath "cosmossdk.io/math" + "github.com/cosmos/cosmos-sdk/codec" + sdk "github.com/cosmos/cosmos-sdk/types" + + "github.com/tokenize-x/tx-chain/v7/x/pse/types" +) + +// MigrateStore migrates the PSE module state from v1 to v2. +// - DelegationTimeEntries key: Pair[AccAddress, ValAddress] -> Triple[uint64, AccAddress, ValAddress]. +// - AccountScoreSnapshot key: AccAddress -> Pair[uint64, AccAddress]. +func MigrateStore( + ctx context.Context, + storeService sdkstore.KVStoreService, + cdc codec.BinaryCodec, +) error { + distributionID, err := getFirstDistributionID(ctx, storeService, cdc) + if err != nil { + return err + } + + if err := migrateDelegationTimeEntries(ctx, storeService, cdc, distributionID); err != nil { + return err + } + + return migrateAccountScoreSnapshot(ctx, storeService, distributionID) +} + +func getFirstDistributionID( + ctx context.Context, + storeService sdkstore.KVStoreService, + cdc codec.BinaryCodec, +) (uint64, error) { + sb := collections.NewSchemaBuilder(storeService) + schedule := collections.NewMap( + sb, + types.AllocationScheduleKey, + "allocation_schedule", + collections.Uint64Key, + codec.CollValue[types.ScheduledDistribution](cdc), + ) + if _, err := sb.Build(); err != nil { + return 0, err + } + + iter, err := schedule.Iterate(ctx, nil) + if err != nil { + return 0, err + } + defer iter.Close() + + if !iter.Valid() { + return 0, nil + } + + kv, err := iter.KeyValue() + if err != nil { + return 0, err + } + + return kv.Value.ID, nil +} + +func migrateDelegationTimeEntries( + ctx context.Context, + storeService sdkstore.KVStoreService, + cdc codec.BinaryCodec, + distributionID uint64, +) error { + oldSB := collections.NewSchemaBuilder(storeService) + oldMap := collections.NewMap( + oldSB, + types.StakingTimeKey, + "delegation_time_entries", + collections.PairKeyCodec(sdk.AccAddressKey, sdk.ValAddressKey), + codec.CollValue[types.DelegationTimeEntry](cdc), + ) + if _, err := oldSB.Build(); err != nil { + return err + } + + type entry struct { + delAddr sdk.AccAddress + valAddr sdk.ValAddress + value types.DelegationTimeEntry + } + + var entries []entry + err := oldMap.Walk(ctx, nil, func( + key collections.Pair[sdk.AccAddress, sdk.ValAddress], + value types.DelegationTimeEntry, + ) (bool, error) { + entries = append(entries, entry{ + delAddr: key.K1(), + valAddr: key.K2(), + value: value, + }) + return false, nil + }) + if err != nil { + return err + } + + if err := oldMap.Clear(ctx, nil); err != nil { + return err + } + + newSB := collections.NewSchemaBuilder(storeService) + newMap := collections.NewMap( + newSB, + types.StakingTimeKey, + "delegation_time_entries", + collections.TripleKeyCodec(collections.Uint64Key, sdk.AccAddressKey, sdk.ValAddressKey), + codec.CollValue[types.DelegationTimeEntry](cdc), + ) + if _, err := newSB.Build(); err != nil { + return err + } + + for _, e := range entries { + key := collections.Join3(distributionID, e.delAddr, e.valAddr) + if err := newMap.Set(ctx, key, e.value); err != nil { + return err + } + } + + return nil +} + +func migrateAccountScoreSnapshot( + ctx context.Context, + storeService sdkstore.KVStoreService, + distributionID uint64, +) error { + oldSB := collections.NewSchemaBuilder(storeService) + oldMap := collections.NewMap( + oldSB, + types.AccountScoreKey, + "account_score", + sdk.AccAddressKey, + sdk.IntValue, + ) + if _, err := oldSB.Build(); err != nil { + return err + } + + type entry struct { + addr sdk.AccAddress + score sdkmath.Int + } + + var entries []entry + err := oldMap.Walk(ctx, nil, func(key sdk.AccAddress, value sdkmath.Int) (bool, error) { + entries = append(entries, entry{addr: key, score: value}) + return false, nil + }) + if err != nil { + return err + } + + if err := oldMap.Clear(ctx, nil); err != nil { + return err + } + + newSB := collections.NewSchemaBuilder(storeService) + newMap := collections.NewMap( + newSB, + types.AccountScoreKey, + "account_score", + collections.PairKeyCodec(collections.Uint64Key, sdk.AccAddressKey), + sdk.IntValue, + ) + if _, err := newSB.Build(); err != nil { + return err + } + + for _, e := range entries { + key := collections.Join(distributionID, e.addr) + if err := newMap.Set(ctx, key, e.score); err != nil { + return err + } + } + + return nil +} diff --git a/x/pse/module.go b/x/pse/module.go index 38aa3b5f..dcc12687 100644 --- a/x/pse/module.go +++ b/x/pse/module.go @@ -100,6 +100,11 @@ func NewAppModule(keeper keeper.Keeper) AppModule { func (am AppModule) RegisterServices(cfg module.Configurator) { types.RegisterMsgServer(cfg.MsgServer(), keeper.NewMsgServer(am.keeper)) types.RegisterQueryServer(cfg.QueryServer(), keeper.NewQueryService(am.keeper)) + + m := keeper.NewMigrator(am.keeper) + if err := cfg.RegisterMigration(types.ModuleName, 1, m.Migrate1to2); err != nil { + panic(err) + } } // Name returns the module's name. @@ -132,7 +137,7 @@ func (am AppModule) IsAppModule() {} func (am AppModule) IsOnePerModuleType() {} // ConsensusVersion implements AppModule/ConsensusVersion. -func (AppModule) ConsensusVersion() uint64 { return 1 } +func (AppModule) ConsensusVersion() uint64 { return 2 } // EndBlock returns the end blocker for the module. It returns no validator // updates. From 0f317625c1bd1546023219ebdfae468ced6b7840 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Mon, 2 Mar 2026 14:23:51 +0300 Subject: [PATCH 18/34] preserve distribution_id in genesis export/import --- x/pse/keeper/genesis.go | 31 +++++++++---------------------- 1 file changed, 9 insertions(+), 22 deletions(-) diff --git a/x/pse/keeper/genesis.go b/x/pse/keeper/genesis.go index 212e4751..70b1b9fe 100644 --- a/x/pse/keeper/genesis.go +++ b/x/pse/keeper/genesis.go @@ -28,36 +28,25 @@ func (k Keeper) InitGenesis(ctx context.Context, genState types.GenesisState) er } } - // TODO revise this logic for distribution id and genesis state - var currentDistributionID uint64 - if len(genState.ScheduledDistributions) > 0 { - currentDistributionID = genState.ScheduledDistributions[0].Timestamp - } - - // Populate delegation time entries from genesis state - for _, delegationTimeEntryExported := range genState.DelegationTimeEntries { - valAddr, err := k.valAddressCodec.StringToBytes(delegationTimeEntryExported.ValidatorAddress) + // Populate delegation time entries from genesis state. + for _, entry := range genState.DelegationTimeEntries { + valAddr, err := k.valAddressCodec.StringToBytes(entry.ValidatorAddress) if err != nil { return err } - delAddr, err := k.addressCodec.StringToBytes(delegationTimeEntryExported.DelegatorAddress) + delAddr, err := k.addressCodec.StringToBytes(entry.DelegatorAddress) if err != nil { return err } - if err = k.SetDelegationTimeEntry( - ctx, - delegationTimeEntryExported.DistributionID, - valAddr, - delAddr, - types.DelegationTimeEntry{ - Shares: delegationTimeEntryExported.Shares, - LastChangedUnixSec: delegationTimeEntryExported.LastChangedUnixSec, - }); err != nil { + if err = k.SetDelegationTimeEntry(ctx, entry.DistributionID, valAddr, delAddr, types.DelegationTimeEntry{ + Shares: entry.Shares, + LastChangedUnixSec: entry.LastChangedUnixSec, + }); err != nil { return err } } - // Populate account scores from genesis state + // Populate account scores from genesis state. for _, accountScore := range genState.AccountScores { addr, err := k.addressCodec.StringToBytes(accountScore.Address) if err != nil { @@ -102,7 +91,6 @@ func (k Keeper) ExportGenesis(ctx context.Context) (*types.GenesisState, error) if err != nil { return false, err } - // TODO revise this logic for distribution id and genesis state delegationTimeEntriesExported = append(delegationTimeEntriesExported, types.DelegationTimeEntryExport{ DistributionID: key.K1(), ValidatorAddress: valAddr, @@ -124,7 +112,6 @@ func (k Keeper) ExportGenesis(ctx context.Context) (*types.GenesisState, error) if err != nil { return false, err } - // TODO revise this logic for distribution id and genesis state genesis.AccountScores = append(genesis.AccountScores, types.AccountScore{ DistributionID: key.K1(), Address: addr, From cb3b06065bcf57590941ada821936963baac99b6 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Tue, 3 Mar 2026 11:06:19 +0300 Subject: [PATCH 19/34] add multi-block EndBlocker routing unit tests --- x/pse/keeper/distribute.go | 2 +- x/pse/keeper/distribution_test.go | 466 ++++++++++++++++++++++++++++++ x/pse/keeper/hooks_test.go | 39 +++ x/pse/migrations/v2/migrate.go | 2 + 4 files changed, 508 insertions(+), 1 deletion(-) diff --git a/x/pse/keeper/distribute.go b/x/pse/keeper/distribute.go index 413af6ed..40ae07dd 100644 --- a/x/pse/keeper/distribute.go +++ b/x/pse/keeper/distribute.go @@ -19,7 +19,7 @@ const defaultBatchSize = 100 // TODO: make configurable // converting each entry into a score snapshot and migrating it to currentID (prevID+1). // // For each entry in the batch: -// 1. Calculate score from lastChanged to distribution timestamp → addToScore(prevID) +// 1. Calculate score from lastChanged to distribution timestamp -> addToScore(prevID) // 2. Create new entry in currentID with same shares, lastChanged = distTimestamp // 3. Remove entry from prevID // diff --git a/x/pse/keeper/distribution_test.go b/x/pse/keeper/distribution_test.go index 65c8cef6..ee786de4 100644 --- a/x/pse/keeper/distribution_test.go +++ b/x/pse/keeper/distribution_test.go @@ -4,9 +4,13 @@ import ( "testing" "time" + "cosmossdk.io/collections" sdkmath "cosmossdk.io/math" "github.com/cosmos/cosmos-sdk/crypto/keys/ed25519" sdk "github.com/cosmos/cosmos-sdk/types" + minttypes "github.com/cosmos/cosmos-sdk/x/mint/types" + stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper" + stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types" "github.com/stretchr/testify/require" "github.com/tokenize-x/tx-chain/v7/testutil/simapp" @@ -266,6 +270,468 @@ func TestDistribution_PrecisionWithMultipleRecipients(t *testing.T) { "community pool should have received the distribution remainders") } +// TestDistribution_MultiBlockEndBlockerRouting tests the full EndBlocker routing logic +// across multiple calls to ProcessNextDistribution, verifying phase transitions: +// +// Call 1 (idle -> start): non-community allocations distributed, OngoingDistribution set +// Call 2 (Phase 1): score conversion batch processed +// Call 3 (Phase 1 -> done): empty batch, TotalScore computed +// Call 4 (Phase 2): tokens distributed to delegators +// Call 5 (Phase 2 -> cleanup): empty batch, cleanup runs, OngoingDistribution removed +// Call 6 (idle): no ongoing, no due schedule, nothing happens +func TestDistribution_MultiBlockEndBlockerRouting(t *testing.T) { + requireT := require.New(t) + + startTime := time.Now().Round(time.Second) + testApp := simapp.New(simapp.WithStartTime(startTime)) + ctx, _, err := testApp.BeginNextBlockAtTime(startTime) + requireT.NoError(err) + + pseKeeper := testApp.PSEKeeper + bankKeeper := testApp.BankKeeper + + bondDenom, err := testApp.StakingKeeper.BondDenom(ctx) + requireT.NoError(err) + + // Create validator + valOp, _ := testApp.GenAccount(ctx) + requireT.NoError(testApp.FundAccount(ctx, valOp, sdk.NewCoins(sdk.NewCoin(bondDenom, sdkmath.NewInt(1000))))) + val, err := testApp.AddValidator(ctx, valOp, sdk.NewInt64Coin(bondDenom, 10), nil) + requireT.NoError(err) + valAddr := sdk.MustValAddressFromBech32(val.GetOperator()) + + // Create two delegators with delegations + del1, _ := testApp.GenAccount(ctx) + del2, _ := testApp.GenAccount(ctx) + requireT.NoError(testApp.FundAccount(ctx, del1, sdk.NewCoins(sdk.NewCoin(bondDenom, sdkmath.NewInt(10_000))))) + requireT.NoError(testApp.FundAccount(ctx, del2, sdk.NewCoins(sdk.NewCoin(bondDenom, sdkmath.NewInt(10_000))))) + + distributionID := uint64(1) + recipientAddr := sdk.AccAddress(ed25519.GenPrivKey().PubKey().Address()).String() + + // Save initial schedule for hooks to find the distribution ID + err = pseKeeper.SaveDistributionSchedule(ctx, []types.ScheduledDistribution{ + {ID: distributionID, Timestamp: distributionID}, + }) + requireT.NoError(err) + + // Delegate + for _, del := range []sdk.AccAddress{del1, del2} { + msg := &stakingtypes.MsgDelegate{ + DelegatorAddress: del.String(), + ValidatorAddress: valAddr.String(), + Amount: sdk.NewInt64Coin(bondDenom, 500), + } + _, err = stakingkeeper.NewMsgServerImpl(testApp.StakingKeeper).Delegate(ctx, msg) + requireT.NoError(err) + } + + // Advance time for score accumulation + ctx, _, err = testApp.BeginNextBlockAtTime(ctx.BlockTime().Add(10 * time.Second)) + requireT.NoError(err) + + // Set up clearing account mappings + params, err := pseKeeper.GetParams(ctx) + requireT.NoError(err) + params.ClearingAccountMappings = []types.ClearingAccountMapping{ + {ClearingAccount: types.ClearingAccountFoundation, RecipientAddresses: []string{recipientAddr}}, + {ClearingAccount: types.ClearingAccountAlliance, RecipientAddresses: []string{recipientAddr}}, + {ClearingAccount: types.ClearingAccountPartnership, RecipientAddresses: []string{recipientAddr}}, + {ClearingAccount: types.ClearingAccountInvestors, RecipientAddresses: []string{recipientAddr}}, + {ClearingAccount: types.ClearingAccountTeam, RecipientAddresses: []string{recipientAddr}}, + } + err = pseKeeper.SetParams(ctx, params) + requireT.NoError(err) + + // Fund all clearing accounts + communityAmount := sdkmath.NewInt(1000) + nonCommunityAmount := sdkmath.NewInt(100) + for _, clearingAccount := range types.GetAllClearingAccounts() { + amount := nonCommunityAmount + if clearingAccount == types.ClearingAccountCommunity { + amount = communityAmount + } + coins := sdk.NewCoins(sdk.NewCoin(bondDenom, amount)) + err = bankKeeper.MintCoins(ctx, minttypes.ModuleName, coins) + requireT.NoError(err) + err = bankKeeper.SendCoinsFromModuleToModule(ctx, minttypes.ModuleName, clearingAccount, coins) + requireT.NoError(err) + } + + // Update schedule with the actual distribution (due now) + distTimestamp := uint64(ctx.BlockTime().Unix()) + err = pseKeeper.AllocationSchedule.Remove(ctx, distributionID) + requireT.NoError(err) + err = pseKeeper.SaveDistributionSchedule(ctx, []types.ScheduledDistribution{ + { + ID: distributionID, + Timestamp: distTimestamp, + Allocations: []types.ClearingAccountAllocation{ + {ClearingAccount: types.ClearingAccountCommunity, Amount: communityAmount}, + {ClearingAccount: types.ClearingAccountFoundation, Amount: nonCommunityAmount}, + {ClearingAccount: types.ClearingAccountAlliance, Amount: nonCommunityAmount}, + {ClearingAccount: types.ClearingAccountPartnership, Amount: nonCommunityAmount}, + {ClearingAccount: types.ClearingAccountInvestors, Amount: nonCommunityAmount}, + {ClearingAccount: types.ClearingAccountTeam, Amount: nonCommunityAmount}, + }, + }, + }) + requireT.NoError(err) + + // --- Call 1: Start distribution --- + err = pseKeeper.ProcessNextDistribution(ctx) + requireT.NoError(err) + + // Verify: OngoingDistribution should be set + ongoing, err := pseKeeper.OngoingDistribution.Get(ctx) + requireT.NoError(err) + requireT.Equal(distributionID, ongoing.ID) + + // Verify: non-community recipient should have received tokens + recipientBalance := bankKeeper.GetBalance(ctx, sdk.MustAccAddressFromBech32(recipientAddr), bondDenom) + requireT.Equal(nonCommunityAmount.MulRaw(5).String(), recipientBalance.Amount.String(), + "recipient should have received all 5 non-community allocations") + + // Verify: TotalScore should NOT exist yet (Phase 1 hasn't run) + _, err = pseKeeper.TotalScore.Get(ctx, distributionID) + requireT.ErrorIs(err, collections.ErrNotFound) + + // --- Call 2: Phase 1 (process score entries) --- + err = pseKeeper.ProcessNextDistribution(ctx) + requireT.NoError(err) + + // TotalScore still not set (entries processed but empty-batch call needed to compute it) + _, err = pseKeeper.TotalScore.Get(ctx, distributionID) + requireT.ErrorIs(err, collections.ErrNotFound) + + // Verify entries migrated from distributionID to distributionID+1 + hasEntries := false + err = pseKeeper.DelegationTimeEntries.Walk(ctx, + collections.NewPrefixedTripleRange[uint64, sdk.AccAddress, sdk.ValAddress](distributionID+1), + func(key collections.Triple[uint64, sdk.AccAddress, sdk.ValAddress], value types.DelegationTimeEntry) (bool, error) { + hasEntries = true + return true, nil + }) + requireT.NoError(err) + requireT.True(hasEntries, "entries should be migrated to next distribution ID") + + // --- Call 3: Phase 1 done (empty batch -> compute TotalScore) --- + err = pseKeeper.ProcessNextDistribution(ctx) + requireT.NoError(err) + + // TotalScore should now exist + totalScore, err := pseKeeper.TotalScore.Get(ctx, distributionID) + requireT.NoError(err) + requireT.True(totalScore.IsPositive(), "TotalScore should be positive") + + // OngoingDistribution should still exist + _, err = pseKeeper.OngoingDistribution.Get(ctx) + requireT.NoError(err) + + // --- Call 4: Phase 2 (distribute tokens) --- + err = pseKeeper.ProcessNextDistribution(ctx) + requireT.NoError(err) + + // OngoingDistribution should still exist (cleanup hasn't run yet) + _, err = pseKeeper.OngoingDistribution.Get(ctx) + requireT.NoError(err) + + // --- Call 5: Phase 2 done (empty batch -> cleanup) --- + err = pseKeeper.ProcessNextDistribution(ctx) + requireT.NoError(err) + + // OngoingDistribution should be removed + _, err = pseKeeper.OngoingDistribution.Get(ctx) + requireT.ErrorIs(err, collections.ErrNotFound, "OngoingDistribution should be removed after cleanup") + + // Schedule entry should be removed + _, err = pseKeeper.AllocationSchedule.Get(ctx, distributionID) + requireT.ErrorIs(err, collections.ErrNotFound, "schedule entry should be removed after cleanup") + + // TotalScore should be cleaned up + _, err = pseKeeper.TotalScore.Get(ctx, distributionID) + requireT.ErrorIs(err, collections.ErrNotFound, "TotalScore should be removed after cleanup") + + // Delegators should have received community tokens (auto-delegated) + stakingQuerier := stakingkeeper.NewQuerier(testApp.StakingKeeper) + for _, del := range []sdk.AccAddress{del1, del2} { + resp, err := stakingQuerier.DelegatorDelegations(ctx, &stakingtypes.QueryDelegatorDelegationsRequest{ + DelegatorAddr: del.String(), + }) + requireT.NoError(err) + totalDelegated := sdkmath.NewInt(0) + for _, d := range resp.DelegationResponses { + totalDelegated = totalDelegated.Add(d.Balance.Amount) + } + requireT.True(totalDelegated.GT(sdkmath.NewInt(500)), + "delegator should have more than initial 500 after community distribution") + } + + // --- Call 6: Idle (nothing to do) --- + err = pseKeeper.ProcessNextDistribution(ctx) + requireT.NoError(err) + + // Still no ongoing + _, err = pseKeeper.OngoingDistribution.Get(ctx) + requireT.ErrorIs(err, collections.ErrNotFound) +} + +// TestDistribution_NonCommunityOnlySingleBlock tests that a distribution with +// no community allocation completes in a single call to ProcessNextDistribution. +func TestDistribution_NonCommunityOnlySingleBlock(t *testing.T) { + requireT := require.New(t) + + testApp := simapp.New() + ctx, _, err := testApp.BeginNextBlock() + requireT.NoError(err) + + pseKeeper := testApp.PSEKeeper + bankKeeper := testApp.BankKeeper + + bondDenom, err := testApp.StakingKeeper.BondDenom(ctx) + requireT.NoError(err) + + recipientAddr := sdk.AccAddress(ed25519.GenPrivKey().PubKey().Address()).String() + + // Set up mappings + params, err := pseKeeper.GetParams(ctx) + requireT.NoError(err) + params.ClearingAccountMappings = []types.ClearingAccountMapping{ + {ClearingAccount: types.ClearingAccountFoundation, RecipientAddresses: []string{recipientAddr}}, + {ClearingAccount: types.ClearingAccountAlliance, RecipientAddresses: []string{recipientAddr}}, + {ClearingAccount: types.ClearingAccountPartnership, RecipientAddresses: []string{recipientAddr}}, + {ClearingAccount: types.ClearingAccountInvestors, RecipientAddresses: []string{recipientAddr}}, + {ClearingAccount: types.ClearingAccountTeam, RecipientAddresses: []string{recipientAddr}}, + } + err = pseKeeper.SetParams(ctx, params) + requireT.NoError(err) + + // Fund non-community clearing accounts only + amount := sdkmath.NewInt(100) + for _, clearingAccount := range types.GetNonCommunityClearingAccounts() { + coins := sdk.NewCoins(sdk.NewCoin(bondDenom, amount)) + err = bankKeeper.MintCoins(ctx, minttypes.ModuleName, coins) + requireT.NoError(err) + err = bankKeeper.SendCoinsFromModuleToModule(ctx, minttypes.ModuleName, clearingAccount, coins) + requireT.NoError(err) + } + + // Schedule with zero community allocation + distTime := uint64(ctx.BlockTime().Unix()) - 1 + err = pseKeeper.SaveDistributionSchedule(ctx, []types.ScheduledDistribution{ + { + ID: 1, + Timestamp: distTime, + Allocations: []types.ClearingAccountAllocation{ + {ClearingAccount: types.ClearingAccountCommunity, Amount: sdkmath.NewInt(0)}, + {ClearingAccount: types.ClearingAccountFoundation, Amount: amount}, + {ClearingAccount: types.ClearingAccountAlliance, Amount: amount}, + {ClearingAccount: types.ClearingAccountPartnership, Amount: amount}, + {ClearingAccount: types.ClearingAccountInvestors, Amount: amount}, + {ClearingAccount: types.ClearingAccountTeam, Amount: amount}, + }, + }, + }) + requireT.NoError(err) + + // Single call should complete everything + err = pseKeeper.ProcessNextDistribution(ctx) + requireT.NoError(err) + + // No OngoingDistribution should be set (no community allocation) + _, err = pseKeeper.OngoingDistribution.Get(ctx) + requireT.ErrorIs(err, collections.ErrNotFound, "no OngoingDistribution for non-community-only distribution") + + // Schedule entry should be removed + _, err = pseKeeper.AllocationSchedule.Get(ctx, 1) + requireT.ErrorIs(err, collections.ErrNotFound, "schedule should be removed after single-block distribution") + + // Recipient should have received all non-community tokens + recipientBalance := bankKeeper.GetBalance(ctx, sdk.MustAccAddressFromBech32(recipientAddr), bondDenom) + requireT.Equal(amount.MulRaw(5).String(), recipientBalance.Amount.String()) +} + +// TestDistribution_EndBlockerWithScenarios mirrors TestKeeper_Distribute scenarios but routes +// through ProcessNextDistribution (the actual EndBlocker entry point) instead of calling +// Phase1/Phase2 directly. This validates the full EndBlocker routing with real delegation flows. +func TestDistribution_EndBlockerWithScenarios(t *testing.T) { + cases := []struct { + name string + actions []func(*runEnv) + }{ + { + name: "unaccumulated score via EndBlocker", + actions: []func(*runEnv){ + func(r *runEnv) { delegateAction(r, r.delegators[0], r.validators[0], 1_100_000) }, + func(r *runEnv) { delegateAction(r, r.delegators[1], r.validators[0], 900_000) }, + func(r *runEnv) { waitAction(r, time.Second*8) }, + func(r *runEnv) { endBlockerDistributeAction(r, sdkmath.NewInt(1000)) }, + func(r *runEnv) { + assertDistributionAction(r, map[*sdk.AccAddress]sdkmath.Int{ + &r.delegators[0]: sdkmath.NewInt(1_100_366), + &r.delegators[1]: sdkmath.NewInt(900_299), + }) + }, + func(r *runEnv) { assertScoreResetAction(r) }, + }, + }, + { + name: "accumulated + unaccumulated score via EndBlocker", + actions: []func(*runEnv){ + func(r *runEnv) { delegateAction(r, r.delegators[0], r.validators[0], 1_100_000) }, + func(r *runEnv) { delegateAction(r, r.delegators[1], r.validators[0], 900_000) }, + func(r *runEnv) { waitAction(r, time.Second*8) }, + func(r *runEnv) { delegateAction(r, r.delegators[0], r.validators[0], 900_000) }, + func(r *runEnv) { delegateAction(r, r.delegators[1], r.validators[0], 1_100_000) }, + func(r *runEnv) { waitAction(r, time.Second*8) }, + func(r *runEnv) { endBlockerDistributeAction(r, sdkmath.NewInt(1000)) }, + func(r *runEnv) { + assertDistributionAction(r, map[*sdk.AccAddress]sdkmath.Int{ + &r.delegators[0]: sdkmath.NewInt(2_000_387), + &r.delegators[1]: sdkmath.NewInt(2_000_362), + }) + }, + func(r *runEnv) { assertScoreResetAction(r) }, + }, + }, + { + name: "unbonding delegation via EndBlocker", + actions: []func(*runEnv){ + func(r *runEnv) { delegateAction(r, r.delegators[0], r.validators[0], 1_100_000) }, + func(r *runEnv) { delegateAction(r, r.delegators[1], r.validators[0], 900_000) }, + func(r *runEnv) { waitAction(r, time.Second*8) }, + func(r *runEnv) { undelegateAction(r, r.delegators[0], r.validators[0], 900_000) }, + func(r *runEnv) { undelegateAction(r, r.delegators[1], r.validators[0], 700_000) }, + func(r *runEnv) { waitAction(r, time.Second*8) }, + func(r *runEnv) { endBlockerDistributeAction(r, sdkmath.NewInt(1000)) }, + func(r *runEnv) { + assertDistributionAction(r, map[*sdk.AccAddress]sdkmath.Int{ + &r.delegators[0]: sdkmath.NewInt(200_295), + &r.delegators[1]: sdkmath.NewInt(200_249), + }) + }, + func(r *runEnv) { assertCommunityPoolBalanceAction(r, sdkmath.NewInt(2)) }, + func(r *runEnv) { assertScoreResetAction(r) }, + }, + }, + { + name: "redelegation via EndBlocker", + actions: []func(*runEnv){ + func(r *runEnv) { delegateAction(r, r.delegators[0], r.validators[0], 1_100_000) }, + func(r *runEnv) { delegateAction(r, r.delegators[1], r.validators[0], 900_000) }, + func(r *runEnv) { waitAction(r, time.Second*8) }, + func(r *runEnv) { redelegateAction(r, r.delegators[0], r.validators[0], r.validators[2], 900_000) }, + func(r *runEnv) { redelegateAction(r, r.delegators[1], r.validators[0], r.validators[2], 700_000) }, + func(r *runEnv) { waitAction(r, time.Second*8) }, + func(r *runEnv) { endBlockerDistributeAction(r, sdkmath.NewInt(1000)) }, + func(r *runEnv) { + assertDistributionAction(r, map[*sdk.AccAddress]sdkmath.Int{ + &r.delegators[0]: sdkmath.NewInt(1_100_365), + &r.delegators[1]: sdkmath.NewInt(900_298), + }) + }, + func(r *runEnv) { assertCommunityPoolBalanceAction(r, sdkmath.NewInt(2)) }, + func(r *runEnv) { assertScoreResetAction(r) }, + }, + }, + { + name: "zero score via EndBlocker", + actions: []func(*runEnv){ + func(r *runEnv) { endBlockerDistributeAction(r, sdkmath.NewInt(1000)) }, + func(r *runEnv) { + assertDistributionAction(r, map[*sdk.AccAddress]sdkmath.Int{ + &r.delegators[0]: sdkmath.NewInt(0), + &r.delegators[1]: sdkmath.NewInt(0), + }) + }, + func(r *runEnv) { assertCommunityPoolBalanceAction(r, sdkmath.NewInt(1000)) }, + func(r *runEnv) { assertScoreResetAction(r) }, + }, + }, + { + name: "multiple distributions via EndBlocker", + actions: []func(*runEnv){ + func(r *runEnv) { delegateAction(r, r.delegators[0], r.validators[0], 1_100_000) }, + func(r *runEnv) { delegateAction(r, r.delegators[1], r.validators[0], 900_000) }, + func(r *runEnv) { waitAction(r, time.Second*8) }, + func(r *runEnv) { endBlockerDistributeAction(r, sdkmath.NewInt(1000)) }, + func(r *runEnv) { + assertDistributionAction(r, map[*sdk.AccAddress]sdkmath.Int{ + &r.delegators[0]: sdkmath.NewInt(1_100_366), + &r.delegators[1]: sdkmath.NewInt(900_299), + }) + }, + func(r *runEnv) { assertCommunityPoolBalanceAction(r, sdkmath.NewInt(2)) }, + func(r *runEnv) { assertScoreResetAction(r) }, + func(r *runEnv) { waitAction(r, time.Second*8) }, + func(r *runEnv) { endBlockerDistributeAction(r, sdkmath.NewInt(1000)) }, + func(r *runEnv) { + assertDistributionAction(r, map[*sdk.AccAddress]sdkmath.Int{ + &r.delegators[0]: sdkmath.NewInt(1_100_732), + &r.delegators[1]: sdkmath.NewInt(900_598), + }) + }, + func(r *runEnv) { assertCommunityPoolBalanceAction(r, sdkmath.NewInt(4)) }, + func(r *runEnv) { assertScoreResetAction(r) }, + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + requireT := require.New(t) + startTime := time.Now().Round(time.Second) + testApp := simapp.New(simapp.WithStartTime(startTime)) + ctx, _, err := testApp.BeginNextBlockAtTime(startTime) + requireT.NoError(err) + runContext := &runEnv{ + testApp: testApp, + ctx: ctx, + requireT: requireT, + currentDistID: tempDistributionID, + } + + // add validators. + for range 3 { + validatorOperator, _ := testApp.GenAccount(ctx) + requireT.NoError(testApp.FundAccount( + ctx, validatorOperator, sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, sdkmath.NewInt(1000)))), + ) + validator, err := testApp.AddValidator( + ctx, validatorOperator, sdk.NewInt64Coin(sdk.DefaultBondDenom, 10), nil, + ) + requireT.NoError(err) + runContext.validators = append( + runContext.validators, + sdk.MustValAddressFromBech32(validator.GetOperator()), + ) + } + + // add delegators. + for range 3 { + delegator, _ := testApp.GenAccount(ctx) + requireT.NoError(testApp.FundAccount( + ctx, delegator, sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, sdkmath.NewInt(1000))), + )) + runContext.delegators = append(runContext.delegators, delegator) + } + + err = testApp.PSEKeeper.SaveDistributionSchedule(ctx, []types.ScheduledDistribution{ + { + Timestamp: tempDistributionID, + ID: tempDistributionID, + }, + }) + requireT.NoError(err) + + // run actions. + for _, action := range tc.actions { + action(runContext) + } + }) + } +} + func TestDistribution_EndBlockFailure(t *testing.T) { requireT := require.New(t) diff --git a/x/pse/keeper/hooks_test.go b/x/pse/keeper/hooks_test.go index 91deb7f7..1fc8a477 100644 --- a/x/pse/keeper/hooks_test.go +++ b/x/pse/keeper/hooks_test.go @@ -1,6 +1,7 @@ package keeper_test import ( + "errors" "testing" "time" @@ -410,6 +411,44 @@ func distributeAction(r *runEnv, amount sdkmath.Int) { r.currentDistID++ } +// endBlockerDistributeAction runs distribution through ProcessNextDistribution (the actual EndBlocker entry point) +// instead of directly calling Phase1/Phase2. This validates the full routing logic. +func endBlockerDistributeAction(r *runEnv, amount sdkmath.Int) { + mintAndSendToPSECommunityClearingAccount(r, amount) + + // Update the AllocationSchedule so ProcessNextDistribution picks it up as due. + scheduledDistribution := types.ScheduledDistribution{ + Timestamp: uint64(r.ctx.BlockTime().Unix()), + ID: r.currentDistID, + Allocations: []types.ClearingAccountAllocation{{ + ClearingAccount: types.ClearingAccountCommunity, + Amount: amount, + }}, + } + err := r.testApp.PSEKeeper.AllocationSchedule.Set(r.ctx, scheduledDistribution.ID, scheduledDistribution) + r.requireT.NoError(err) + + // Call ProcessNextDistribution repeatedly until distribution completes. + for i := range 20 { + err = r.testApp.PSEKeeper.ProcessNextDistribution(r.ctx) + r.requireT.NoError(err, "ProcessNextDistribution failed at iteration %d", i) + + _, ongoingErr := r.testApp.PSEKeeper.OngoingDistribution.Get(r.ctx) + if errors.Is(ongoingErr, collections.ErrNotFound) { + // Cleanup done — distribution complete. + break + } + r.requireT.NoError(ongoingErr) + } + + // Verify cleanup completed. + _, err = r.testApp.PSEKeeper.OngoingDistribution.Get(r.ctx) + r.requireT.ErrorIs(err, collections.ErrNotFound, "OngoingDistribution should be removed after distribution") + + // Advance to next distribution ID (Phase 1 migrated entries to currentDistID+1). + r.currentDistID++ +} + func mintAndSendCoin(r *runEnv, recipient sdk.AccAddress, coins sdk.Coins) { r.requireT.NoError( r.testApp.BankKeeper.MintCoins(r.ctx, minttypes.ModuleName, coins), diff --git a/x/pse/migrations/v2/migrate.go b/x/pse/migrations/v2/migrate.go index e4548e5d..fee57e94 100644 --- a/x/pse/migrations/v2/migrate.go +++ b/x/pse/migrations/v2/migrate.go @@ -32,6 +32,8 @@ func MigrateStore( return migrateAccountScoreSnapshot(ctx, storeService, distributionID) } +// TODO: Currently assigns the first distribution ID to all entries. Implement proper mapping +// of entries to correct distribution IDs based on timestamps when multiple distributions exist. func getFirstDistributionID( ctx context.Context, storeService sdkstore.KVStoreService, From 5645f15b446bf50033de5da2ec67cb9eb1089423 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Tue, 3 Mar 2026 13:14:47 +0300 Subject: [PATCH 20/34] fix typo --- x/pse/keeper/distribute_test.go | 1 - 1 file changed, 1 deletion(-) diff --git a/x/pse/keeper/distribute_test.go b/x/pse/keeper/distribute_test.go index 2c3d5fcf..ae5b74ca 100644 --- a/x/pse/keeper/distribute_test.go +++ b/x/pse/keeper/distribute_test.go @@ -350,7 +350,6 @@ func Test_ExcludedAddress_FullLifecycle(t *testing.T) { scheduledDistribution := types.ScheduledDistribution{ ID: distributionID, Timestamp: uint64(ctx.BlockTime().Unix()), - ID: distributionID, Allocations: []types.ClearingAccountAllocation{{ ClearingAccount: types.ClearingAccountCommunity, Amount: amount, From dd535a159327cd427265cf87963cd86c2ddca3de Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Wed, 4 Mar 2026 18:27:53 +0300 Subject: [PATCH 21/34] reorder formatting for PeekNextAllocationSchedule --- x/pse/keeper/distribution.go | 64 ++++++++++++++++++------------------ 1 file changed, 32 insertions(+), 32 deletions(-) diff --git a/x/pse/keeper/distribution.go b/x/pse/keeper/distribution.go index e6c59fa7..82e71396 100644 --- a/x/pse/keeper/distribution.go +++ b/x/pse/keeper/distribution.go @@ -120,6 +120,38 @@ func (k Keeper) resumeOngoingDistribution(ctx context.Context, ongoing types.Sch return nil } +// PeekNextAllocationSchedule returns the earliest scheduled distribution and whether it should be processed. +func (k Keeper) PeekNextAllocationSchedule(ctx context.Context) (types.ScheduledDistribution, bool, error) { + sdkCtx := sdk.UnwrapSDKContext(ctx) + + // Get iterator for the allocation schedule (sorted by id ascending) + iter, err := k.AllocationSchedule.Iterate(ctx, nil) + if err != nil { + return types.ScheduledDistribution{}, false, err + } + defer iter.Close() + + // Return early if schedule is empty + if !iter.Valid() { + return types.ScheduledDistribution{}, false, nil + } + + // Extract the earliest scheduled distribution (sorted by id ascending) + kv, err := iter.KeyValue() + if err != nil { + return types.ScheduledDistribution{}, false, err + } + + scheduledDist := kv.Value + + // Check if distribution time has arrived + // Since IDs are sequential and timestamps are monotonically increasing, + // the first item by ID is also the earliest by time. + shouldProcess := scheduledDist.Timestamp <= uint64(sdkCtx.BlockTime().Unix()) + + return scheduledDist, shouldProcess, nil +} + // distributeNonCommunityAllocations processes all non-community allocations in a single block. func (k Keeper) distributeNonCommunityAllocations( ctx context.Context, @@ -206,38 +238,6 @@ func (k Keeper) distributeNonCommunityAllocations( return nil } -// PeekNextAllocationSchedule returns the earliest scheduled distribution and whether it should be processed. -func (k Keeper) PeekNextAllocationSchedule(ctx context.Context) (types.ScheduledDistribution, bool, error) { - sdkCtx := sdk.UnwrapSDKContext(ctx) - - // Get iterator for the allocation schedule (sorted by id ascending) - iter, err := k.AllocationSchedule.Iterate(ctx, nil) - if err != nil { - return types.ScheduledDistribution{}, false, err - } - defer iter.Close() - - // Return early if schedule is empty - if !iter.Valid() { - return types.ScheduledDistribution{}, false, nil - } - - // Extract the earliest scheduled distribution (sorted by id ascending) - kv, err := iter.KeyValue() - if err != nil { - return types.ScheduledDistribution{}, false, err - } - - scheduledDist := kv.Value - - // Check if distribution time has arrived - // Since IDs are sequential and timestamps are monotonically increasing, - // the first item by ID is also the earliest by time. - shouldProcess := scheduledDist.Timestamp <= uint64(sdkCtx.BlockTime().Unix()) - - return scheduledDist, shouldProcess, nil -} - // SaveDistributionSchedule persists the distribution schedule to blockchain state. // Each scheduled distribution is stored in the AllocationSchedule map, indexed by its ID. func (k Keeper) SaveDistributionSchedule(ctx context.Context, schedule []types.ScheduledDistribution) error { From 3681d87f27150313cb47de578dc9660c52ef58c8 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Thu, 5 Mar 2026 14:32:04 +0300 Subject: [PATCH 22/34] fix: remove prevID entry in migrateOngoingEntry to prevent double scoring --- x/pse/keeper/hooks.go | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/x/pse/keeper/hooks.go b/x/pse/keeper/hooks.go index 369b5922..aaeed47e 100644 --- a/x/pse/keeper/hooks.go +++ b/x/pse/keeper/hooks.go @@ -150,13 +150,9 @@ func (h Hooks) BeforeDelegationRemoved(ctx context.Context, delAddr sdk.AccAddre return err } if ongoingFound { - handled, err := h.migrateOngoingEntry(ctx, ongoing, currentID, delAddr, valAddr, blockTime) - if err != nil { + if _, err := h.migrateOngoingEntry(ctx, ongoing, currentID, delAddr, valAddr, blockTime); err != nil { return err } - if handled { - return h.k.RemoveDelegationTimeEntry(ctx, ongoing.ID, valAddr, delAddr) - } } // Scenario 2: Entry exists in current distribution. @@ -304,5 +300,10 @@ func (h Hooks) migrateOngoingEntry( return false, err } + // Remove the old entry from prevID to prevent double scoring in Phase 1 batch processing. + if err := h.k.RemoveDelegationTimeEntry(ctx, prevID, valAddr, delAddr); err != nil { + return false, err + } + return true, nil } From cca3d25ab495c76ab7a85acc571264bfaf17dd14 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Thu, 5 Mar 2026 14:43:07 +0300 Subject: [PATCH 23/34] rename prevID/currentID to ongoingID/nextID for clarity --- x/pse/keeper/distribute.go | 50 ++++++++++++------------- x/pse/keeper/distribution.go | 8 ++-- x/pse/keeper/hooks.go | 72 ++++++++++++++++++------------------ x/pse/keeper/hooks_test.go | 6 +-- 4 files changed, 68 insertions(+), 68 deletions(-) diff --git a/x/pse/keeper/distribute.go b/x/pse/keeper/distribute.go index 40ae07dd..fc53edbb 100644 --- a/x/pse/keeper/distribute.go +++ b/x/pse/keeper/distribute.go @@ -15,24 +15,24 @@ import ( // defaultBatchSize is the number of entries processed per EndBlock during multi-block distribution. const defaultBatchSize = 100 // TODO: make configurable -// ProcessPhase1ScoreConversion processes a batch of DelegationTimeEntries from the ongoing distribution (prevID), -// converting each entry into a score snapshot and migrating it to currentID (prevID+1). +// ProcessPhase1ScoreConversion processes a batch of DelegationTimeEntries from the ongoing distribution (ongoingID), +// converting each entry into a score snapshot and migrating it to nextID (ongoingID + 1). // // For each entry in the batch: -// 1. Calculate score from lastChanged to distribution timestamp -> addToScore(prevID) -// 2. Create new entry in currentID with same shares, lastChanged = distTimestamp -// 3. Remove entry from prevID +// 1. Calculate score from lastChanged to distribution timestamp -> addToScore(ongoingID) +// 2. Create new entry under nextID with same shares, lastChanged = distTimestamp +// 3. Remove entry from ongoingID // -// Returns true when all prevID entries have been processed and TotalScore is computed. +// Returns true when all ongoingID entries have been processed and TotalScore is computed. func (k Keeper) ProcessPhase1ScoreConversion(ctx context.Context, ongoing types.ScheduledDistribution) (bool, error) { - prevID := ongoing.ID - currentID := ongoing.ID + 1 + ongoingID := ongoing.ID + nextID := ongoing.ID + 1 distTimestamp := int64(ongoing.Timestamp) - // Collect a batch of entries from prevID. + // Collect a batch of entries from ongoingID. iter, err := k.DelegationTimeEntries.Iterate( ctx, - collections.NewPrefixedTripleRange[uint64, sdk.AccAddress, sdk.ValAddress](prevID), + collections.NewPrefixedTripleRange[uint64, sdk.AccAddress, sdk.ValAddress](ongoingID), ) if err != nil { return false, err @@ -61,7 +61,7 @@ func (k Keeper) ProcessPhase1ScoreConversion(ctx context.Context, ongoing types. // Compute TotalScore from all accumulated snapshots. if len(batch) == 0 { - if err := k.computeTotalScore(ctx, prevID); err != nil { + if err := k.computeTotalScore(ctx, ongoingID); err != nil { return false, err } return true, nil @@ -78,21 +78,21 @@ func (k Keeper) ProcessPhase1ScoreConversion(ctx context.Context, ongoing types. if err != nil { return false, err } - if err := k.addToScore(ctx, prevID, item.delAddr, score); err != nil { + if err := k.addToScore(ctx, ongoingID, item.delAddr, score); err != nil { return false, err } } - // Migrate entry to currentID with same shares, reset lastChanged to distribution timestamp. - if err := k.SetDelegationTimeEntry(ctx, currentID, item.valAddr, item.delAddr, types.DelegationTimeEntry{ + // Migrate entry to nextID with same shares, reset lastChanged to distribution timestamp. + if err := k.SetDelegationTimeEntry(ctx, nextID, item.valAddr, item.delAddr, types.DelegationTimeEntry{ LastChangedUnixSec: distTimestamp, Shares: item.entry.Shares, }); err != nil { return false, err } - // Remove from prevID. - if err := k.RemoveDelegationTimeEntry(ctx, prevID, item.valAddr, item.delAddr); err != nil { + // Remove from ongoingID. + if err := k.RemoveDelegationTimeEntry(ctx, ongoingID, item.valAddr, item.delAddr); err != nil { return false, err } } @@ -124,7 +124,7 @@ func (k Keeper) computeTotalScore(ctx context.Context, distributionID uint64) er } // ProcessPhase2TokenDistribution distributes tokens to delegators in batches based on their computed scores. -// Uses TotalScore[prevID] for proportion calculation and iterates AccountScoreSnapshot[prevID]. +// Uses TotalScore[ongoingID] for proportion calculation and iterates AccountScoreSnapshot[ongoingID]. // // For each delegator in the batch: // 1. Compute share: userAmount = totalPSEAmount × score / totalScore @@ -137,10 +137,10 @@ func (k Keeper) computeTotalScore(ctx context.Context, distributionID uint64) er func (k Keeper) ProcessPhase2TokenDistribution( ctx context.Context, ongoing types.ScheduledDistribution, bondDenom string, ) (bool, error) { - prevID := ongoing.ID + ongoingID := ongoing.ID totalPSEAmount := getCommunityAllocationAmount(ongoing) - totalScore, err := k.TotalScore.Get(ctx, prevID) + totalScore, err := k.TotalScore.Get(ctx, ongoingID) if err != nil { return false, err } @@ -152,13 +152,13 @@ func (k Keeper) ProcessPhase2TokenDistribution( return false, err } } - return true, k.cleanupDistribution(ctx, prevID) + return true, k.cleanupDistribution(ctx, ongoingID) } // Collect a batch of score snapshots. iter, err := k.AccountScoreSnapshot.Iterate( ctx, - collections.NewPrefixedPairRange[uint64, sdk.AccAddress](prevID), + collections.NewPrefixedPairRange[uint64, sdk.AccAddress](ongoingID), ) if err != nil { return false, err @@ -186,7 +186,7 @@ func (k Keeper) ProcessPhase2TokenDistribution( // Only triggered when all distributions of this round are completed. // Send leftover to community pool and clean up. if len(batch) == 0 { - distributedSoFar, err := k.getDistributedAmount(ctx, prevID) + distributedSoFar, err := k.getDistributedAmount(ctx, ongoingID) if err != nil { return false, err } @@ -196,7 +196,7 @@ func (k Keeper) ProcessPhase2TokenDistribution( return false, err } } - return true, k.cleanupDistribution(ctx, prevID) + return true, k.cleanupDistribution(ctx, ongoingID) } sdkCtx := sdk.UnwrapSDKContext(ctx) @@ -209,7 +209,7 @@ func (k Keeper) ProcessPhase2TokenDistribution( return false, err } - if err := k.addToDistributedAmount(ctx, prevID, distributedAmount); err != nil { + if err := k.addToDistributedAmount(ctx, ongoingID, distributedAmount); err != nil { return false, err } @@ -224,7 +224,7 @@ func (k Keeper) ProcessPhase2TokenDistribution( } // Remove processed snapshot. - if err := k.RemoveDelegatorScore(ctx, prevID, item.delAddr); err != nil { + if err := k.RemoveDelegatorScore(ctx, ongoingID, item.delAddr); err != nil { return false, err } } diff --git a/x/pse/keeper/distribution.go b/x/pse/keeper/distribution.go index 82e71396..d9203e52 100644 --- a/x/pse/keeper/distribution.go +++ b/x/pse/keeper/distribution.go @@ -85,10 +85,10 @@ func (k Keeper) ProcessNextDistribution(ctx context.Context) error { // Phase is determined by TotalScore existence: absent -> Phase 1, present -> Phase 2. func (k Keeper) resumeOngoingDistribution(ctx context.Context, ongoing types.ScheduledDistribution) error { sdkCtx := sdk.UnwrapSDKContext(ctx) - prevID := ongoing.ID + ongoingID := ongoing.ID // TotalScore absent -> Phase 1 (score conversion still in progress). - _, err := k.TotalScore.Get(ctx, prevID) + _, err := k.TotalScore.Get(ctx, ongoingID) if errors.Is(err, collections.ErrNotFound) { done, err := k.ProcessPhase1ScoreConversion(ctx, ongoing) if err != nil { @@ -96,7 +96,7 @@ func (k Keeper) resumeOngoingDistribution(ctx context.Context, ongoing types.Sch } if done { sdkCtx.Logger().Info("phase 1 complete, TotalScore computed", - "distribution_id", prevID) + "distribution_id", ongoingID) } return nil } @@ -115,7 +115,7 @@ func (k Keeper) resumeOngoingDistribution(ctx context.Context, ongoing types.Sch } if done { sdkCtx.Logger().Info("multi-block community distribution complete", - "distribution_id", prevID) + "distribution_id", ongoingID) } return nil } diff --git a/x/pse/keeper/hooks.go b/x/pse/keeper/hooks.go index aaeed47e..e5d5a124 100644 --- a/x/pse/keeper/hooks.go +++ b/x/pse/keeper/hooks.go @@ -36,11 +36,11 @@ func (k Keeper) getOngoingDistribution(ctx context.Context) (types.ScheduledDist return ongoing, true, nil } -// getCurrentDistributionID returns the distribution ID that new entries should be written to. -// If an ongoing distribution exists (ID=N is being processed), returns N+1. +// getNextDistributionID returns the next distribution ID that new entries should be written to. +// If an ongoing distribution exists (ongoingID=N is being processed), returns N+1. // Otherwise returns the next scheduled distribution's ID (zero-value ID when no schedule exists). // TODO: handle empty distribution schedule — currently returns 0 when no schedule exists. -func (k Keeper) getCurrentDistributionID(ctx context.Context) (uint64, error) { +func (k Keeper) getNextDistributionID(ctx context.Context) (uint64, error) { ongoing, found, err := k.getOngoingDistribution(ctx) if err != nil { return 0, err @@ -58,8 +58,8 @@ func (k Keeper) getCurrentDistributionID(ctx context.Context) (uint64, error) { // AfterDelegationModified implements the staking hooks interface. // Handles 3 scenarios based on where the delegator's entry exists: -// - Scenario 1: Entry in prevID (ongoing distribution in progress). -// - Scenario 2: Entry in currentID — normal score calculation. +// - Scenario 1: Entry in ongoingID (ongoing distribution in progress). +// - Scenario 2: Entry in nextID — normal score calculation. // - Scenario 3: No entry — create new entry, no score. func (h Hooks) AfterDelegationModified(ctx context.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress) error { delegation, err := h.k.stakingKeeper.GetDelegation(ctx, delAddr, valAddr) @@ -67,7 +67,7 @@ func (h Hooks) AfterDelegationModified(ctx context.Context, delAddr sdk.AccAddre return err } - currentID, err := h.k.getCurrentDistributionID(ctx) + nextID, err := h.k.getNextDistributionID(ctx) if err != nil { return err } @@ -83,35 +83,35 @@ func (h Hooks) AfterDelegationModified(ctx context.Context, delAddr sdk.AccAddre blockTime := sdk.UnwrapSDKContext(ctx).BlockTime().Unix() // Scenario 1: Entry exists in previous distribution (ongoing distribution in progress). - // Split score at distribution timestamp, move entry to currentID. + // Split score at distribution timestamp, move entry to nextID. ongoing, ongoingFound, err := h.k.getOngoingDistribution(ctx) if err != nil { return err } if ongoingFound { - handled, err := h.migrateOngoingEntry(ctx, ongoing, currentID, delAddr, valAddr, blockTime) + handled, err := h.migrateOngoingEntry(ctx, ongoing, nextID, delAddr, valAddr, blockTime) if err != nil { return err } if handled { - return h.k.SetDelegationTimeEntry(ctx, currentID, valAddr, delAddr, types.DelegationTimeEntry{ + return h.k.SetDelegationTimeEntry(ctx, nextID, valAddr, delAddr, types.DelegationTimeEntry{ LastChangedUnixSec: blockTime, Shares: delegation.Shares, }) } } - // Scenario 2: Entry exists in current distribution. - currentEntry, err := h.k.GetDelegationTimeEntry(ctx, currentID, valAddr, delAddr) + // Scenario 2: Entry exists in next distribution. + currentEntry, err := h.k.GetDelegationTimeEntry(ctx, nextID, valAddr, delAddr) if err == nil { score, err := calculateAddedScore(ctx, h.k, valAddr, currentEntry) if err != nil { return err } - if err := h.k.addToScore(ctx, currentID, delAddr, score); err != nil { + if err := h.k.addToScore(ctx, nextID, delAddr, score); err != nil { return err } - return h.k.SetDelegationTimeEntry(ctx, currentID, valAddr, delAddr, types.DelegationTimeEntry{ + return h.k.SetDelegationTimeEntry(ctx, nextID, valAddr, delAddr, types.DelegationTimeEntry{ LastChangedUnixSec: blockTime, Shares: delegation.Shares, }) @@ -120,8 +120,8 @@ func (h Hooks) AfterDelegationModified(ctx context.Context, delAddr sdk.AccAddre return err } - // Scenario 3: No entry - create new in currentID (no score, duration = 0). - return h.k.SetDelegationTimeEntry(ctx, currentID, valAddr, delAddr, types.DelegationTimeEntry{ + // Scenario 3: No entry — create new under nextID (no score, duration = 0). + return h.k.SetDelegationTimeEntry(ctx, nextID, valAddr, delAddr, types.DelegationTimeEntry{ LastChangedUnixSec: blockTime, Shares: delegation.Shares, }) @@ -129,7 +129,7 @@ func (h Hooks) AfterDelegationModified(ctx context.Context, delAddr sdk.AccAddre // BeforeDelegationRemoved implements the staking hooks interface. func (h Hooks) BeforeDelegationRemoved(ctx context.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress) error { - currentID, err := h.k.getCurrentDistributionID(ctx) + nextID, err := h.k.getNextDistributionID(ctx) if err != nil { return err } @@ -150,22 +150,22 @@ func (h Hooks) BeforeDelegationRemoved(ctx context.Context, delAddr sdk.AccAddre return err } if ongoingFound { - if _, err := h.migrateOngoingEntry(ctx, ongoing, currentID, delAddr, valAddr, blockTime); err != nil { + if _, err := h.migrateOngoingEntry(ctx, ongoing, nextID, delAddr, valAddr, blockTime); err != nil { return err } } - // Scenario 2: Entry exists in current distribution. - currentEntry, err := h.k.GetDelegationTimeEntry(ctx, currentID, valAddr, delAddr) + // Scenario 2: Entry exists in next distribution. + currentEntry, err := h.k.GetDelegationTimeEntry(ctx, nextID, valAddr, delAddr) if err == nil { score, err := calculateAddedScore(ctx, h.k, valAddr, currentEntry) if err != nil { return err } - if err := h.k.addToScore(ctx, currentID, delAddr, score); err != nil { + if err := h.k.addToScore(ctx, nextID, delAddr, score); err != nil { return err } - return h.k.RemoveDelegationTimeEntry(ctx, currentID, valAddr, delAddr) + return h.k.RemoveDelegationTimeEntry(ctx, nextID, valAddr, delAddr) } if !errors.Is(err, collections.ErrNotFound) { return err @@ -256,19 +256,19 @@ func (h Hooks) AfterUnbondingInitiated(_ context.Context, _ uint64) error { return nil } -// migrateOngoingEntry handles a delegation entry that still lives in the previous (ongoing) distribution. -// It calculates score for both the prev and current periods, removes the entry from prevID, +// migrateOngoingEntry handles a delegation entry that still lives under the ongoing distribution. +// It calculates score for both the ongoing and next periods, removes the entry from ongoingID, // and returns true if the entry was found and processed. func (h Hooks) migrateOngoingEntry( ctx context.Context, ongoing types.ScheduledDistribution, - currentID uint64, + nextID uint64, delAddr sdk.AccAddress, valAddr sdk.ValAddress, blockTime int64, ) (bool, error) { - prevID := ongoing.ID - prevEntry, err := h.k.GetDelegationTimeEntry(ctx, prevID, valAddr, delAddr) + ongoingID := ongoing.ID + ongoingEntry, err := h.k.GetDelegationTimeEntry(ctx, ongoingID, valAddr, delAddr) if errors.Is(err, collections.ErrNotFound) { return false, nil } @@ -278,30 +278,30 @@ func (h Hooks) migrateOngoingEntry( distTimestamp := int64(ongoing.Timestamp) - // Score for previous period: lastChanged -> distribution timestamp. - prevScore, err := calculateScoreAtTimestamp(ctx, h.k, valAddr, prevEntry, distTimestamp) + // Score for ongoing period: lastChanged -> distribution timestamp. + ongoingScore, err := calculateScoreAtTimestamp(ctx, h.k, valAddr, ongoingEntry, distTimestamp) if err != nil { return false, err } - if err := h.k.addToScore(ctx, prevID, delAddr, prevScore); err != nil { + if err := h.k.addToScore(ctx, ongoingID, delAddr, ongoingScore); err != nil { return false, err } - // Score for current period: distribution timestamp -> now. - currentPeriodEntry := types.DelegationTimeEntry{ + // Score for next period: distribution timestamp -> now. + nextPeriodEntry := types.DelegationTimeEntry{ LastChangedUnixSec: distTimestamp, - Shares: prevEntry.Shares, + Shares: ongoingEntry.Shares, } - currentScore, err := calculateScoreAtTimestamp(ctx, h.k, valAddr, currentPeriodEntry, blockTime) + nextScore, err := calculateScoreAtTimestamp(ctx, h.k, valAddr, nextPeriodEntry, blockTime) if err != nil { return false, err } - if err := h.k.addToScore(ctx, currentID, delAddr, currentScore); err != nil { + if err := h.k.addToScore(ctx, nextID, delAddr, nextScore); err != nil { return false, err } - // Remove the old entry from prevID to prevent double scoring in Phase 1 batch processing. - if err := h.k.RemoveDelegationTimeEntry(ctx, prevID, valAddr, delAddr); err != nil { + // Remove the old entry from ongoingID to prevent double scoring in Phase 1 batch processing. + if err := h.k.RemoveDelegationTimeEntry(ctx, ongoingID, valAddr, delAddr); err != nil { return false, err } diff --git a/x/pse/keeper/hooks_test.go b/x/pse/keeper/hooks_test.go index 1fc8a477..ed8a4875 100644 --- a/x/pse/keeper/hooks_test.go +++ b/x/pse/keeper/hooks_test.go @@ -288,9 +288,9 @@ func assertCommunityPoolBalanceAction(r *runEnv, expectedBalance sdkmath.Int) { } func assertScoreResetAction(r *runEnv) { - // After cleanup, score snapshots at the previous distribution ID should be cleared. - prevID := r.currentDistID - 1 - scoreRanger := collections.NewPrefixedPairRange[uint64, sdk.AccAddress](prevID) + // After cleanup, score snapshots at the ongoing distribution ID should be cleared. + ongoingID := r.currentDistID - 1 + scoreRanger := collections.NewPrefixedPairRange[uint64, sdk.AccAddress](ongoingID) err := r.testApp.PSEKeeper.AccountScoreSnapshot.Walk(r.ctx, scoreRanger, func(key collections.Pair[uint64, sdk.AccAddress], value sdkmath.Int) (bool, error) { r.requireT.Equal(sdkmath.NewInt(0), value) From bed8a12ea3d3bf1b7aee5e94f13bdaafb0b44761 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Fri, 6 Mar 2026 09:42:15 +0300 Subject: [PATCH 24/34] apply division by zero bugfix --- x/pse/keeper/distribute.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x/pse/keeper/distribute.go b/x/pse/keeper/distribute.go index fc53edbb..8985ff43 100644 --- a/x/pse/keeper/distribute.go +++ b/x/pse/keeper/distribute.go @@ -313,7 +313,7 @@ func (k Keeper) distributeToDelegator( totalDelegationAmount = totalDelegationAmount.Add(delegation.Balance.Amount) } - if len(delegations) == 0 { + if len(delegations) == 0 || totalDelegationAmount.IsZero() { return sdkmath.NewInt(0), nil } From 8e9498e7f7b133be144bf646ae813ec1a4a48909 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Mon, 9 Mar 2026 12:42:50 +0300 Subject: [PATCH 25/34] fix PSE fairness: honor earned rewards regardless of undelegations & only skip auto-delegation --- x/pse/keeper/distribute.go | 13 +++++++++---- x/pse/keeper/distribute_test.go | 13 +++++++++---- 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/x/pse/keeper/distribute.go b/x/pse/keeper/distribute.go index 8985ff43..8990d63c 100644 --- a/x/pse/keeper/distribute.go +++ b/x/pse/keeper/distribute.go @@ -313,10 +313,8 @@ func (k Keeper) distributeToDelegator( totalDelegationAmount = totalDelegationAmount.Add(delegation.Balance.Amount) } - if len(delegations) == 0 || totalDelegationAmount.IsZero() { - return sdkmath.NewInt(0), nil - } - + // Send earned tokens to delegator's wallet regardless of active delegations. + // Score was earned by staking during the scoring period, so the reward is always honored. if err = k.bankKeeper.SendCoinsFromModuleToAccount( ctx, types.ClearingAccountCommunity, @@ -325,6 +323,13 @@ func (k Keeper) distributeToDelegator( ); err != nil { return sdkmath.NewInt(0), err } + + // Auto-delegate proportionally to active validators. If no active delegations + // (e.g., user undelegated during multi-block distribution), skip auto-delegation + if len(delegations) == 0 || totalDelegationAmount.IsZero() { + return amount, nil + } + for _, delegation := range delegations { // NOTE: this division will have rounding errors up to 1 subunit, which is acceptable and will be ignored. // if that one subunit exists, it will remain in user balance as undelegated. diff --git a/x/pse/keeper/distribute_test.go b/x/pse/keeper/distribute_test.go index ae5b74ca..509d9335 100644 --- a/x/pse/keeper/distribute_test.go +++ b/x/pse/keeper/distribute_test.go @@ -132,13 +132,18 @@ func TestKeeper_Distribute(t *testing.T) { func(r *runEnv) { undelegateAction(r, r.delegators[0], r.validators[0], 1_100_000) }, func(r *runEnv) { distributeAction(r, sdkmath.NewInt(1000)) }, func(r *runEnv) { + // delegators[0] fully undelegated — no auto-delegation, but earned reward sent as liquid tokens assertDistributionAction(r, map[*sdk.AccAddress]sdkmath.Int{ - &r.delegators[0]: sdkmath.NewInt(0), // + 1000 * 1.1 / 3 - &r.delegators[1]: sdkmath.NewInt(900_299), // + 1000 * 0.9 / 3 + &r.delegators[0]: sdkmath.NewInt(0), // staking balance 0 (no active delegation for auto-delegate) + &r.delegators[1]: sdkmath.NewInt(900_299), // 900k original + 1000 * 0.9 / 2.4 ≈ 299 auto-delegated }) + // delegators[0] receives 366 as liquid: 1000 (FundAccount) + 366 (PSE reward) = 1366 + // undelegated 1,100,000 tokens are in unbonding queue, not liquid + balance := r.testApp.BankKeeper.GetBalance(r.ctx, r.delegators[0], sdk.DefaultBondDenom) + r.requireT.Equal(sdkmath.NewInt(1366), balance.Amount) }, - // + 1000 * 1.1 / 3 (from user's share) + 2 (from rounding) - func(r *runEnv) { assertCommunityPoolBalanceAction(r, sdkmath.NewInt(366+2)) }, + // only rounding leftover goes to community pool (no forfeited rewards) + func(r *runEnv) { assertCommunityPoolBalanceAction(r, sdkmath.NewInt(2)) }, func(r *runEnv) { assertScoreResetAction(r) }, }, }, From 24d36d2231c2ede42f4fbba651b2792d8df7f549 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Thu, 12 Mar 2026 10:34:51 +0300 Subject: [PATCH 26/34] rename Phase1/Phase2 functions per PR review feedback --- x/pse/keeper/distribute.go | 16 ++++++++-------- x/pse/keeper/distribute_test.go | 4 ++-- x/pse/keeper/distribution.go | 4 ++-- x/pse/keeper/hooks_test.go | 4 ++-- 4 files changed, 14 insertions(+), 14 deletions(-) diff --git a/x/pse/keeper/distribute.go b/x/pse/keeper/distribute.go index 8990d63c..c2473711 100644 --- a/x/pse/keeper/distribute.go +++ b/x/pse/keeper/distribute.go @@ -15,7 +15,7 @@ import ( // defaultBatchSize is the number of entries processed per EndBlock during multi-block distribution. const defaultBatchSize = 100 // TODO: make configurable -// ProcessPhase1ScoreConversion processes a batch of DelegationTimeEntries from the ongoing distribution (ongoingID), +// ConsumeOngoingDelegationTimeEntry processes a batch of DelegationTimeEntries from the ongoing distribution (ongoingID), // converting each entry into a score snapshot and migrating it to nextID (ongoingID + 1). // // For each entry in the batch: @@ -24,7 +24,7 @@ const defaultBatchSize = 100 // TODO: make configurable // 3. Remove entry from ongoingID // // Returns true when all ongoingID entries have been processed and TotalScore is computed. -func (k Keeper) ProcessPhase1ScoreConversion(ctx context.Context, ongoing types.ScheduledDistribution) (bool, error) { +func (k Keeper) ConsumeOngoingDelegationTimeEntry(ctx context.Context, ongoing types.ScheduledDistribution) (bool, error) { ongoingID := ongoing.ID nextID := ongoing.ID + 1 distTimestamp := int64(ongoing.Timestamp) @@ -123,7 +123,7 @@ func (k Keeper) computeTotalScore(ctx context.Context, distributionID uint64) er return k.TotalScore.Set(ctx, distributionID, totalScore) } -// ProcessPhase2TokenDistribution distributes tokens to delegators in batches based on their computed scores. +// ProcessOngoingTokenDistribution distributes tokens to delegators in batches based on their computed scores. // Uses TotalScore[ongoingID] for proportion calculation and iterates AccountScoreSnapshot[ongoingID]. // // For each delegator in the batch: @@ -134,7 +134,7 @@ func (k Keeper) computeTotalScore(ctx context.Context, distributionID uint64) er // // When all delegators have been processed, sends leftover (rounding errors + undelegated users) to the community pool. // Returns true when distribution is complete and all state has been cleaned up. -func (k Keeper) ProcessPhase2TokenDistribution( +func (k Keeper) ProcessOngoingTokenDistribution( ctx context.Context, ongoing types.ScheduledDistribution, bondDenom string, ) (bool, error) { ongoingID := ongoing.ID @@ -152,7 +152,7 @@ func (k Keeper) ProcessPhase2TokenDistribution( return false, err } } - return true, k.cleanupDistribution(ctx, ongoingID) + return true, k.cleanupOngoingDistribution(ctx, ongoingID) } // Collect a batch of score snapshots. @@ -196,7 +196,7 @@ func (k Keeper) ProcessPhase2TokenDistribution( return false, err } } - return true, k.cleanupDistribution(ctx, ongoingID) + return true, k.cleanupOngoingDistribution(ctx, ongoingID) } sdkCtx := sdk.UnwrapSDKContext(ctx) @@ -248,8 +248,8 @@ func (k Keeper) sendLeftoverToCommunityPool(ctx context.Context, amount sdkmath. return k.distributionKeeper.FundCommunityPool(ctx, sdk.NewCoins(sdk.NewCoin(bondDenom, amount)), pseModuleAddress) } -// cleanupDistribution removes all state associated with a completed distribution. -func (k Keeper) cleanupDistribution(ctx context.Context, distributionID uint64) error { +// cleanupOngoingDistribution removes all state associated with a completed distribution. +func (k Keeper) cleanupOngoingDistribution(ctx context.Context, distributionID uint64) error { if err := k.AccountScoreSnapshot.Clear( ctx, collections.NewPrefixedPairRange[uint64, sdk.AccAddress](distributionID), diff --git a/x/pse/keeper/distribute_test.go b/x/pse/keeper/distribute_test.go index 509d9335..2386ea5a 100644 --- a/x/pse/keeper/distribute_test.go +++ b/x/pse/keeper/distribute_test.go @@ -364,14 +364,14 @@ func Test_ExcludedAddress_FullLifecycle(t *testing.T) { requireT.NoError(err) balanceBefore := testApp.BankKeeper.GetBalance(ctx, delAddr, bondDenom) for { - done, err := pseKeeper.ProcessPhase1ScoreConversion(ctx, scheduledDistribution) + done, err := pseKeeper.ConsumeOngoingDelegationTimeEntry(ctx, scheduledDistribution) requireT.NoError(err) if done { break } } for { - done, err := pseKeeper.ProcessPhase2TokenDistribution(ctx, scheduledDistribution, bondDenom) + done, err := pseKeeper.ProcessOngoingTokenDistribution(ctx, scheduledDistribution, bondDenom) requireT.NoError(err) if done { break diff --git a/x/pse/keeper/distribution.go b/x/pse/keeper/distribution.go index d9203e52..d21f1958 100644 --- a/x/pse/keeper/distribution.go +++ b/x/pse/keeper/distribution.go @@ -90,7 +90,7 @@ func (k Keeper) resumeOngoingDistribution(ctx context.Context, ongoing types.Sch // TotalScore absent -> Phase 1 (score conversion still in progress). _, err := k.TotalScore.Get(ctx, ongoingID) if errors.Is(err, collections.ErrNotFound) { - done, err := k.ProcessPhase1ScoreConversion(ctx, ongoing) + done, err := k.ConsumeOngoingDelegationTimeEntry(ctx, ongoing) if err != nil { return err } @@ -109,7 +109,7 @@ func (k Keeper) resumeOngoingDistribution(ctx context.Context, ongoing types.Sch if err != nil { return err } - done, err := k.ProcessPhase2TokenDistribution(ctx, ongoing, bondDenom) + done, err := k.ProcessOngoingTokenDistribution(ctx, ongoing, bondDenom) if err != nil { return err } diff --git a/x/pse/keeper/hooks_test.go b/x/pse/keeper/hooks_test.go index ed8a4875..4653d73e 100644 --- a/x/pse/keeper/hooks_test.go +++ b/x/pse/keeper/hooks_test.go @@ -391,7 +391,7 @@ func distributeAction(r *runEnv, amount sdkmath.Int) { // Run Phase 1 until done. for { - done, err := r.testApp.PSEKeeper.ProcessPhase1ScoreConversion(r.ctx, scheduledDistribution) + done, err := r.testApp.PSEKeeper.ConsumeOngoingDelegationTimeEntry(r.ctx, scheduledDistribution) r.requireT.NoError(err) if done { break @@ -400,7 +400,7 @@ func distributeAction(r *runEnv, amount sdkmath.Int) { // Run Phase 2 until done. for { - done, err := r.testApp.PSEKeeper.ProcessPhase2TokenDistribution(r.ctx, scheduledDistribution, bondDenom) + done, err := r.testApp.PSEKeeper.ProcessOngoingTokenDistribution(r.ctx, scheduledDistribution, bondDenom) r.requireT.NoError(err) if done { break From da84216597e46a0f02e6e03f4a20da49cbfa18eb Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Thu, 12 Mar 2026 10:42:46 +0300 Subject: [PATCH 27/34] revert fairness fix: send undelegated users' tokens to community pool per review --- x/pse/keeper/distribute.go | 13 +++++-------- x/pse/keeper/distribute_test.go | 13 ++++++------- 2 files changed, 11 insertions(+), 15 deletions(-) diff --git a/x/pse/keeper/distribute.go b/x/pse/keeper/distribute.go index c2473711..126b59b7 100644 --- a/x/pse/keeper/distribute.go +++ b/x/pse/keeper/distribute.go @@ -313,8 +313,11 @@ func (k Keeper) distributeToDelegator( totalDelegationAmount = totalDelegationAmount.Add(delegation.Balance.Amount) } - // Send earned tokens to delegator's wallet regardless of active delegations. - // Score was earned by staking during the scoring period, so the reward is always honored. + // Only distribute to users with active stakes. If not, it will be leftover. + if len(delegations) == 0 || totalDelegationAmount.IsZero() { + return sdkmath.NewInt(0), nil + } + if err = k.bankKeeper.SendCoinsFromModuleToAccount( ctx, types.ClearingAccountCommunity, @@ -324,12 +327,6 @@ func (k Keeper) distributeToDelegator( return sdkmath.NewInt(0), err } - // Auto-delegate proportionally to active validators. If no active delegations - // (e.g., user undelegated during multi-block distribution), skip auto-delegation - if len(delegations) == 0 || totalDelegationAmount.IsZero() { - return amount, nil - } - for _, delegation := range delegations { // NOTE: this division will have rounding errors up to 1 subunit, which is acceptable and will be ignored. // if that one subunit exists, it will remain in user balance as undelegated. diff --git a/x/pse/keeper/distribute_test.go b/x/pse/keeper/distribute_test.go index 2386ea5a..fcb733fa 100644 --- a/x/pse/keeper/distribute_test.go +++ b/x/pse/keeper/distribute_test.go @@ -132,18 +132,17 @@ func TestKeeper_Distribute(t *testing.T) { func(r *runEnv) { undelegateAction(r, r.delegators[0], r.validators[0], 1_100_000) }, func(r *runEnv) { distributeAction(r, sdkmath.NewInt(1000)) }, func(r *runEnv) { - // delegators[0] fully undelegated — no auto-delegation, but earned reward sent as liquid tokens + // delegators[0] fully undelegated — no active delegations, tokens go to community pool assertDistributionAction(r, map[*sdk.AccAddress]sdkmath.Int{ - &r.delegators[0]: sdkmath.NewInt(0), // staking balance 0 (no active delegation for auto-delegate) + &r.delegators[0]: sdkmath.NewInt(0), // no active delegation -> tokens go to community pool &r.delegators[1]: sdkmath.NewInt(900_299), // 900k original + 1000 * 0.9 / 2.4 ≈ 299 auto-delegated }) - // delegators[0] receives 366 as liquid: 1000 (FundAccount) + 366 (PSE reward) = 1366 - // undelegated 1,100,000 tokens are in unbonding queue, not liquid + // delegators[0] only has original 1000 funded amount (no PSE reward) balance := r.testApp.BankKeeper.GetBalance(r.ctx, r.delegators[0], sdk.DefaultBondDenom) - r.requireT.Equal(sdkmath.NewInt(1366), balance.Amount) + r.requireT.Equal(sdkmath.NewInt(1000), balance.Amount) }, - // only rounding leftover goes to community pool (no forfeited rewards) - func(r *runEnv) { assertCommunityPoolBalanceAction(r, sdkmath.NewInt(2)) }, + // delegators[0]'s share (366) + rounding (2) goes to community pool + func(r *runEnv) { assertCommunityPoolBalanceAction(r, sdkmath.NewInt(368)) }, func(r *runEnv) { assertScoreResetAction(r) }, }, }, From 3525a9dd412fe99c076fe837274c2881ae559f73 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Thu, 12 Mar 2026 11:18:25 +0300 Subject: [PATCH 28/34] accumulate TotalScore incrementally via addToScore --- x/pse/keeper/delegation.go | 16 ++++++++++++-- x/pse/keeper/distribute.go | 36 ++++++------------------------- x/pse/keeper/distribution.go | 24 ++++++--------------- x/pse/keeper/distribution_test.go | 36 ++++++++----------------------- 4 files changed, 36 insertions(+), 76 deletions(-) diff --git a/x/pse/keeper/delegation.go b/x/pse/keeper/delegation.go index 7350fb3e..fe6581da 100644 --- a/x/pse/keeper/delegation.go +++ b/x/pse/keeper/delegation.go @@ -69,7 +69,8 @@ func (k Keeper) RemoveDelegatorScore(ctx context.Context, distributionID uint64, return k.AccountScoreSnapshot.Remove(ctx, key) } -// addToScore atomically adds a score value to a delegator's score snapshot. +// addToScore atomically adds a score value to a delegator's score snapshot +// and incrementally updates TotalScore for the same distribution. func (k Keeper) addToScore( ctx context.Context, distributionID uint64, delAddr sdk.AccAddress, score sdkmath.Int, ) error { @@ -82,7 +83,18 @@ func (k Keeper) addToScore( } else if err != nil { return err } - return k.SetDelegatorScore(ctx, distributionID, delAddr, lastScore.Add(score)) + if err := k.SetDelegatorScore(ctx, distributionID, delAddr, lastScore.Add(score)); err != nil { + return err + } + + // Accumulate TotalScore + currentTotal, err := k.TotalScore.Get(ctx, distributionID) + if errors.Is(err, collections.ErrNotFound) { + currentTotal = sdkmath.NewInt(0) + } else if err != nil { + return err + } + return k.TotalScore.Set(ctx, distributionID, currentTotal.Add(score)) } // CalculateDelegatorScore calculates the current total score for a delegator. diff --git a/x/pse/keeper/distribute.go b/x/pse/keeper/distribute.go index 126b59b7..b2d94910 100644 --- a/x/pse/keeper/distribute.go +++ b/x/pse/keeper/distribute.go @@ -23,7 +23,7 @@ const defaultBatchSize = 100 // TODO: make configurable // 2. Create new entry under nextID with same shares, lastChanged = distTimestamp // 3. Remove entry from ongoingID // -// Returns true when all ongoingID entries have been processed and TotalScore is computed. +// Returns true when all ongoingID entries have been processed. func (k Keeper) ConsumeOngoingDelegationTimeEntry(ctx context.Context, ongoing types.ScheduledDistribution) (bool, error) { ongoingID := ongoing.ID nextID := ongoing.ID + 1 @@ -59,11 +59,7 @@ func (k Keeper) ConsumeOngoingDelegationTimeEntry(ctx context.Context, ongoing t } iter.Close() - // Compute TotalScore from all accumulated snapshots. if len(batch) == 0 { - if err := k.computeTotalScore(ctx, ongoingID); err != nil { - return false, err - } return true, nil } @@ -97,30 +93,8 @@ func (k Keeper) ConsumeOngoingDelegationTimeEntry(ctx context.Context, ongoing t } } - return false, nil -} - -// computeTotalScore sums all AccountScoreSnapshot entries for a distribution and stores the result in TotalScore. -func (k Keeper) computeTotalScore(ctx context.Context, distributionID uint64) error { - iter, err := k.AccountScoreSnapshot.Iterate( - ctx, - collections.NewPrefixedPairRange[uint64, sdk.AccAddress](distributionID), - ) - if err != nil { - return err - } - defer iter.Close() - - totalScore := sdkmath.NewInt(0) - for ; iter.Valid(); iter.Next() { - kv, err := iter.KeyValue() - if err != nil { - return err - } - totalScore = totalScore.Add(kv.Value) - } - - return k.TotalScore.Set(ctx, distributionID, totalScore) + // If batch was smaller than the limit, all entries have been consumed. + return len(batch) < defaultBatchSize, nil } // ProcessOngoingTokenDistribution distributes tokens to delegators in batches based on their computed scores. @@ -141,7 +115,9 @@ func (k Keeper) ProcessOngoingTokenDistribution( totalPSEAmount := getCommunityAllocationAmount(ongoing) totalScore, err := k.TotalScore.Get(ctx, ongoingID) - if err != nil { + if errors.Is(err, collections.ErrNotFound) { + totalScore = sdkmath.NewInt(0) + } else if err != nil { return false, err } diff --git a/x/pse/keeper/distribution.go b/x/pse/keeper/distribution.go index d21f1958..84010166 100644 --- a/x/pse/keeper/distribution.go +++ b/x/pse/keeper/distribution.go @@ -2,9 +2,7 @@ package keeper import ( "context" - "errors" - "cosmossdk.io/collections" errorsmod "cosmossdk.io/errors" sdkmath "cosmossdk.io/math" sdk "github.com/cosmos/cosmos-sdk/types" @@ -82,29 +80,21 @@ func (k Keeper) ProcessNextDistribution(ctx context.Context) error { } // resumeOngoingDistribution continues a multi-block community distribution. -// Phase is determined by TotalScore existence: absent -> Phase 1, present -> Phase 2. +// Consumes DelegationTimeEntries in batches, then distributes tokens once all entries are consumed. func (k Keeper) resumeOngoingDistribution(ctx context.Context, ongoing types.ScheduledDistribution) error { sdkCtx := sdk.UnwrapSDKContext(ctx) ongoingID := ongoing.ID - // TotalScore absent -> Phase 1 (score conversion still in progress). - _, err := k.TotalScore.Get(ctx, ongoingID) - if errors.Is(err, collections.ErrNotFound) { - done, err := k.ConsumeOngoingDelegationTimeEntry(ctx, ongoing) - if err != nil { - return err - } - if done { - sdkCtx.Logger().Info("phase 1 complete, TotalScore computed", - "distribution_id", ongoingID) - } - return nil - } + // Consume remaining DelegationTimeEntries for score conversion. + isConsumed, err := k.ConsumeOngoingDelegationTimeEntry(ctx, ongoing) if err != nil { return err } + if !isConsumed { + return nil + } - // TotalScore present -> Phase 2 (token distribution). + // All entries consumed — distribute tokens. bondDenom, err := k.stakingKeeper.BondDenom(ctx) if err != nil { return err diff --git a/x/pse/keeper/distribution_test.go b/x/pse/keeper/distribution_test.go index ee786de4..6c4e1477 100644 --- a/x/pse/keeper/distribution_test.go +++ b/x/pse/keeper/distribution_test.go @@ -392,17 +392,16 @@ func TestDistribution_MultiBlockEndBlockerRouting(t *testing.T) { requireT.Equal(nonCommunityAmount.MulRaw(5).String(), recipientBalance.Amount.String(), "recipient should have received all 5 non-community allocations") - // Verify: TotalScore should NOT exist yet (Phase 1 hasn't run) - _, err = pseKeeper.TotalScore.Get(ctx, distributionID) - requireT.ErrorIs(err, collections.ErrNotFound) - - // --- Call 2: Phase 1 (process score entries) --- + // --- Call 2: Consume all entries + distribute tokens --- + // Batch size (100) > delegator count (2), so all entries are consumed in one batch (isConsumed=true). + // Tokens distributed to all delegators in one batch. err = pseKeeper.ProcessNextDistribution(ctx) requireT.NoError(err) - // TotalScore still not set (entries processed but empty-batch call needed to compute it) - _, err = pseKeeper.TotalScore.Get(ctx, distributionID) - requireT.ErrorIs(err, collections.ErrNotFound) + // TotalScore is accumulated incrementally via addToScore. + totalScore, err := pseKeeper.TotalScore.Get(ctx, distributionID) + requireT.NoError(err) + requireT.True(totalScore.IsPositive(), "TotalScore should be positive") // Verify entries migrated from distributionID to distributionID+1 hasEntries := false @@ -415,28 +414,11 @@ func TestDistribution_MultiBlockEndBlockerRouting(t *testing.T) { requireT.NoError(err) requireT.True(hasEntries, "entries should be migrated to next distribution ID") - // --- Call 3: Phase 1 done (empty batch -> compute TotalScore) --- - err = pseKeeper.ProcessNextDistribution(ctx) - requireT.NoError(err) - - // TotalScore should now exist - totalScore, err := pseKeeper.TotalScore.Get(ctx, distributionID) - requireT.NoError(err) - requireT.True(totalScore.IsPositive(), "TotalScore should be positive") - - // OngoingDistribution should still exist - _, err = pseKeeper.OngoingDistribution.Get(ctx) - requireT.NoError(err) - - // --- Call 4: Phase 2 (distribute tokens) --- - err = pseKeeper.ProcessNextDistribution(ctx) - requireT.NoError(err) - - // OngoingDistribution should still exist (cleanup hasn't run yet) + // OngoingDistribution should still exist (cleanup runs on next empty-batch call) _, err = pseKeeper.OngoingDistribution.Get(ctx) requireT.NoError(err) - // --- Call 5: Phase 2 done (empty batch -> cleanup) --- + // --- Call 3: Cleanup (no entries, no snapshots -> cleanup) --- err = pseKeeper.ProcessNextDistribution(ctx) requireT.NoError(err) From 63f311f1f2733104399f9b6adc9b559f5ca34db6 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Thu, 12 Mar 2026 11:54:08 +0300 Subject: [PATCH 29/34] add comments to distributeNonCommunityAllocations --- x/pse/keeper/distribution.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/x/pse/keeper/distribution.go b/x/pse/keeper/distribution.go index 84010166..5d0b1361 100644 --- a/x/pse/keeper/distribution.go +++ b/x/pse/keeper/distribution.go @@ -161,6 +161,7 @@ func (k Keeper) distributeNonCommunityAllocations( continue } + // Look up recipient addresses for this clearing account from governance-configured mappings. var recipientAddrs []string for _, mapping := range clearingAccountMappings { if mapping.ClearingAccount == allocation.ClearingAccount { @@ -169,6 +170,7 @@ func (k Keeper) distributeNonCommunityAllocations( } } + // Split allocation evenly among recipients; remainder goes to community pool. numRecipients := sdkmath.NewInt(int64(len(recipientAddrs))) if numRecipients.IsZero() { return errorsmod.Wrapf( @@ -180,6 +182,7 @@ func (k Keeper) distributeNonCommunityAllocations( amountPerRecipient := allocation.Amount.Quo(numRecipients) remainder := allocation.Amount.Mod(numRecipients) + // Send equal share to each recipient from the clearing account. for _, recipientAddr := range recipientAddrs { recipient := sdk.MustAccAddressFromBech32(recipientAddr) coinsToSend := sdk.NewCoins(sdk.NewCoin(bondDenom, amountPerRecipient)) @@ -200,6 +203,7 @@ func (k Keeper) distributeNonCommunityAllocations( } } + // Send remainder to the community pool. if !remainder.IsZero() { clearingAccountAddr := k.accountKeeper.GetModuleAddress(allocation.ClearingAccount) remainderCoins := sdk.NewCoins(sdk.NewCoin(bondDenom, remainder)) From f9ff57ba770e3ebd17c5e531d90ecd2b414ed4f7 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Thu, 12 Mar 2026 12:24:28 +0300 Subject: [PATCH 30/34] add phase-1 gap score for nextID to ensure fairness across migration batches --- x/pse/keeper/distribute.go | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/x/pse/keeper/distribute.go b/x/pse/keeper/distribute.go index b2d94910..686d74c6 100644 --- a/x/pse/keeper/distribute.go +++ b/x/pse/keeper/distribute.go @@ -20,14 +20,17 @@ const defaultBatchSize = 100 // TODO: make configurable // // For each entry in the batch: // 1. Calculate score from lastChanged to distribution timestamp -> addToScore(ongoingID) -// 2. Create new entry under nextID with same shares, lastChanged = distTimestamp -// 3. Remove entry from ongoingID +// 2. Calculate gap score from distribution timestamp to current block time -> addToScore(nextID) +// 3. Create new entry under nextID with same shares, lastChanged = current block time +// 4. Remove entry from ongoingID // // Returns true when all ongoingID entries have been processed. func (k Keeper) ConsumeOngoingDelegationTimeEntry(ctx context.Context, ongoing types.ScheduledDistribution) (bool, error) { + sdkCtx := sdk.UnwrapSDKContext(ctx) ongoingID := ongoing.ID nextID := ongoing.ID + 1 distTimestamp := int64(ongoing.Timestamp) + blockTime := sdkCtx.BlockTime().Unix() // Collect a batch of entries from ongoingID. iter, err := k.DelegationTimeEntries.Iterate( @@ -70,6 +73,7 @@ func (k Keeper) ConsumeOngoingDelegationTimeEntry(ctx context.Context, ongoing t } if !isExcluded { + // Score for ongoingID: lastChanged -> distTimestamp. score, err := calculateScoreAtTimestamp(ctx, k, item.valAddr, item.entry, distTimestamp) if err != nil { return false, err @@ -77,11 +81,24 @@ func (k Keeper) ConsumeOngoingDelegationTimeEntry(ctx context.Context, ongoing t if err := k.addToScore(ctx, ongoingID, item.delAddr, score); err != nil { return false, err } + + // Score for nextID: distTimestamp -> blockTime (gap during Phase 1 processing). + // TODO: add dedicated integration test to verify gap score fairness across batches. + gapScore, err := calculateScoreAtTimestamp(ctx, k, item.valAddr, types.DelegationTimeEntry{ + LastChangedUnixSec: distTimestamp, + Shares: item.entry.Shares, + }, blockTime) + if err != nil { + return false, err + } + if err := k.addToScore(ctx, nextID, item.delAddr, gapScore); err != nil { + return false, err + } } - // Migrate entry to nextID with same shares, reset lastChanged to distribution timestamp. + // Migrate entry to nextID with same shares, lastChanged = current block time. if err := k.SetDelegationTimeEntry(ctx, nextID, item.valAddr, item.delAddr, types.DelegationTimeEntry{ - LastChangedUnixSec: distTimestamp, + LastChangedUnixSec: blockTime, Shares: item.entry.Shares, }); err != nil { return false, err From 2e44d215081214d8f2698912d632a01f96c2c0ad Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Thu, 12 Mar 2026 12:27:57 +0300 Subject: [PATCH 31/34] lint fix --- x/pse/keeper/distribute.go | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/x/pse/keeper/distribute.go b/x/pse/keeper/distribute.go index 686d74c6..4b06443d 100644 --- a/x/pse/keeper/distribute.go +++ b/x/pse/keeper/distribute.go @@ -15,8 +15,9 @@ import ( // defaultBatchSize is the number of entries processed per EndBlock during multi-block distribution. const defaultBatchSize = 100 // TODO: make configurable -// ConsumeOngoingDelegationTimeEntry processes a batch of DelegationTimeEntries from the ongoing distribution (ongoingID), -// converting each entry into a score snapshot and migrating it to nextID (ongoingID + 1). +// ConsumeOngoingDelegationTimeEntry processes a batch of DelegationTimeEntries +// from the ongoing distribution (ongoingID), converting each entry into a score +// snapshot and migrating it to nextID (ongoingID + 1). // // For each entry in the batch: // 1. Calculate score from lastChanged to distribution timestamp -> addToScore(ongoingID) @@ -25,7 +26,9 @@ const defaultBatchSize = 100 // TODO: make configurable // 4. Remove entry from ongoingID // // Returns true when all ongoingID entries have been processed. -func (k Keeper) ConsumeOngoingDelegationTimeEntry(ctx context.Context, ongoing types.ScheduledDistribution) (bool, error) { +func (k Keeper) ConsumeOngoingDelegationTimeEntry( + ctx context.Context, ongoing types.ScheduledDistribution, +) (bool, error) { sdkCtx := sdk.UnwrapSDKContext(ctx) ongoingID := ongoing.ID nextID := ongoing.ID + 1 From 328e6391ce477c42672ed508ceb9413b2a35e1fc Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Thu, 12 Mar 2026 12:53:29 +0300 Subject: [PATCH 32/34] invariant: disable pse module for positive PSE amunt but non-positive score --- x/pse/keeper/distribute.go | 20 ++++++++---- x/pse/keeper/distribute_test.go | 10 ++---- x/pse/keeper/distribution_test.go | 49 ++++++++++++++++++++-------- x/pse/keeper/hooks_test.go | 53 +++++++++++++++++++++++++++++++ x/pse/types/errors.go | 3 ++ 5 files changed, 106 insertions(+), 29 deletions(-) diff --git a/x/pse/keeper/distribute.go b/x/pse/keeper/distribute.go index 4b06443d..8b8d6abf 100644 --- a/x/pse/keeper/distribute.go +++ b/x/pse/keeper/distribute.go @@ -5,6 +5,7 @@ import ( "errors" "cosmossdk.io/collections" + errorsmod "cosmossdk.io/errors" sdkmath "cosmossdk.io/math" sdk "github.com/cosmos/cosmos-sdk/types" stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types" @@ -141,13 +142,18 @@ func (k Keeper) ProcessOngoingTokenDistribution( return false, err } - // No score or no amount: send everything to community pool and clean up. - if !totalScore.IsPositive() || totalPSEAmount.IsZero() { - if totalPSEAmount.IsPositive() { - if err := k.sendLeftoverToCommunityPool(ctx, totalPSEAmount, bondDenom); err != nil { - return false, err - } - } + // Invariant: positive amount with non-positive score indicates a scoring bug. + // Return error and disable PSE. + if totalPSEAmount.IsPositive() && !totalScore.IsPositive() { + return false, errorsmod.Wrapf( + types.ErrInvariantViolation, + "positive PSE amount %s but non-positive total score %s for distribution %d", + totalPSEAmount, totalScore, ongoingID, + ) + } + + // No amount to distribute: clean up. + if !totalPSEAmount.IsPositive() { return true, k.cleanupOngoingDistribution(ctx, ongoingID) } diff --git a/x/pse/keeper/distribute_test.go b/x/pse/keeper/distribute_test.go index fcb733fa..4a65d75f 100644 --- a/x/pse/keeper/distribute_test.go +++ b/x/pse/keeper/distribute_test.go @@ -147,17 +147,11 @@ func TestKeeper_Distribute(t *testing.T) { }, }, { - name: "zero score", + name: "zero score triggers invariant violation", actions: []func(*runEnv){ - func(r *runEnv) { distributeAction(r, sdkmath.NewInt(1000)) }, func(r *runEnv) { - assertDistributionAction(r, map[*sdk.AccAddress]sdkmath.Int{ - &r.delegators[0]: sdkmath.NewInt(0), - &r.delegators[1]: sdkmath.NewInt(0), - }) + distributeExpectInvariantViolation(r, sdkmath.NewInt(1000)) }, - func(r *runEnv) { assertCommunityPoolBalanceAction(r, sdkmath.NewInt(1000)) }, - func(r *runEnv) { assertScoreResetAction(r) }, }, }, { diff --git a/x/pse/keeper/distribution_test.go b/x/pse/keeper/distribution_test.go index 6c4e1477..e66ba4aa 100644 --- a/x/pse/keeper/distribution_test.go +++ b/x/pse/keeper/distribution_test.go @@ -20,15 +20,45 @@ import ( func TestDistribution_GenesisRebuild(t *testing.T) { requireT := require.New(t) - testApp := simapp.New() - ctx := testApp.NewContext(false) - ctx = ctx.WithBlockTime(time.Now()) // Set proper block time + startTime := time.Now().Round(time.Second) + testApp := simapp.New(simapp.WithStartTime(startTime)) + ctx, _, err := testApp.BeginNextBlockAtTime(startTime) + requireT.NoError(err) pseKeeper := testApp.PSEKeeper // Get bond denom bondDenom, err := testApp.StakingKeeper.BondDenom(ctx) requireT.NoError(err) + // Create a validator and delegator so community distribution has non-zero score. + valOp, _ := testApp.GenAccount(ctx) + requireT.NoError(testApp.FundAccount(ctx, valOp, sdk.NewCoins(sdk.NewCoin(bondDenom, sdkmath.NewInt(1000))))) + val, errVal := testApp.AddValidator(ctx, valOp, sdk.NewInt64Coin(bondDenom, 10), nil) + requireT.NoError(errVal) + valAddr := sdk.MustValAddressFromBech32(val.GetOperator()) + + del1, _ := testApp.GenAccount(ctx) + requireT.NoError(testApp.FundAccount(ctx, del1, sdk.NewCoins(sdk.NewCoin(bondDenom, sdkmath.NewInt(10_000))))) + + time1 := uint64(startTime.Add(1 * time.Hour).Unix()) + time2 := uint64(startTime.Add(2 * time.Hour).Unix()) + + // Save initial schedule so hooks can find distribution ID. + err = pseKeeper.SaveDistributionSchedule(ctx, []types.ScheduledDistribution{ + {ID: 1, Timestamp: time1}, + }) + requireT.NoError(err) + + // Delegate and advance time for score accumulation. + _, err = stakingkeeper.NewMsgServerImpl(testApp.StakingKeeper).Delegate(ctx, &stakingtypes.MsgDelegate{ + DelegatorAddress: del1.String(), + ValidatorAddress: valAddr.String(), + Amount: sdk.NewInt64Coin(bondDenom, 500), + }) + requireT.NoError(err) + ctx, _, err = testApp.BeginNextBlockAtTime(ctx.BlockTime().Add(10 * time.Second)) + requireT.NoError(err) + // Set up mappings and fund modules for all eligible accounts addr1 := sdk.AccAddress(ed25519.GenPrivKey().PubKey().Address()).String() addr2 := sdk.AccAddress(ed25519.GenPrivKey().PubKey().Address()).String() @@ -54,9 +84,6 @@ func TestDistribution_GenesisRebuild(t *testing.T) { requireT.NoError(err) } - time1 := uint64(time.Now().Add(1 * time.Hour).Unix()) - time2 := uint64(time.Now().Add(2 * time.Hour).Unix()) - // Set up params with mappings params, err := pseKeeper.GetParams(ctx) requireT.NoError(err) @@ -617,17 +644,11 @@ func TestDistribution_EndBlockerWithScenarios(t *testing.T) { }, }, { - name: "zero score via EndBlocker", + name: "zero score via EndBlocker triggers invariant violation", actions: []func(*runEnv){ - func(r *runEnv) { endBlockerDistributeAction(r, sdkmath.NewInt(1000)) }, func(r *runEnv) { - assertDistributionAction(r, map[*sdk.AccAddress]sdkmath.Int{ - &r.delegators[0]: sdkmath.NewInt(0), - &r.delegators[1]: sdkmath.NewInt(0), - }) + endBlockerDistributeExpectInvariantViolation(r, sdkmath.NewInt(1000)) }, - func(r *runEnv) { assertCommunityPoolBalanceAction(r, sdkmath.NewInt(1000)) }, - func(r *runEnv) { assertScoreResetAction(r) }, }, }, { diff --git a/x/pse/keeper/hooks_test.go b/x/pse/keeper/hooks_test.go index 4653d73e..9952adeb 100644 --- a/x/pse/keeper/hooks_test.go +++ b/x/pse/keeper/hooks_test.go @@ -449,6 +449,59 @@ func endBlockerDistributeAction(r *runEnv, amount sdkmath.Int) { r.currentDistID++ } +// distributeExpectInvariantViolation runs Phase 1 + Phase 2 and expects +// ErrInvariantViolation from Phase 2 (e.g., zero score with positive pse amount). +func distributeExpectInvariantViolation(r *runEnv, amount sdkmath.Int) { + mintAndSendToPSECommunityClearingAccount(r, amount) + bondDenom, err := r.testApp.StakingKeeper.BondDenom(r.ctx) + r.requireT.NoError(err) + scheduledDistribution := types.ScheduledDistribution{ + Timestamp: uint64(r.ctx.BlockTime().Unix()), + ID: r.currentDistID, + Allocations: []types.ClearingAccountAllocation{{ + ClearingAccount: types.ClearingAccountCommunity, + Amount: amount, + }}, + } + err = r.testApp.PSEKeeper.OngoingDistribution.Set(r.ctx, scheduledDistribution) + r.requireT.NoError(err) + + for { + done, err := r.testApp.PSEKeeper.ConsumeOngoingDelegationTimeEntry(r.ctx, scheduledDistribution) + r.requireT.NoError(err) + if done { + break + } + } + + _, err = r.testApp.PSEKeeper.ProcessOngoingTokenDistribution(r.ctx, scheduledDistribution, bondDenom) + r.requireT.ErrorIs(err, types.ErrInvariantViolation) +} + +// endBlockerDistributeExpectInvariantViolation runs distribution via +// ProcessNextDistribution and expects ErrInvariantViolation. +func endBlockerDistributeExpectInvariantViolation(r *runEnv, amount sdkmath.Int) { + mintAndSendToPSECommunityClearingAccount(r, amount) + scheduledDistribution := types.ScheduledDistribution{ + Timestamp: uint64(r.ctx.BlockTime().Unix()), + ID: r.currentDistID, + Allocations: []types.ClearingAccountAllocation{{ + ClearingAccount: types.ClearingAccountCommunity, + Amount: amount, + }}, + } + err := r.testApp.PSEKeeper.AllocationSchedule.Set(r.ctx, scheduledDistribution.ID, scheduledDistribution) + r.requireT.NoError(err) + + // First call starts the distribution (sets OngoingDistribution). + err = r.testApp.PSEKeeper.ProcessNextDistribution(r.ctx) + r.requireT.NoError(err) + + // Second call: Phase 1 completes + Phase 2 hits invariant violation. + err = r.testApp.PSEKeeper.ProcessNextDistribution(r.ctx) + r.requireT.ErrorIs(err, types.ErrInvariantViolation) +} + func mintAndSendCoin(r *runEnv, recipient sdk.AccAddress, coins sdk.Coins) { r.requireT.NoError( r.testApp.BankKeeper.MintCoins(r.ctx, minttypes.ModuleName, coins), diff --git a/x/pse/types/errors.go b/x/pse/types/errors.go index d279f870..c6cc877e 100644 --- a/x/pse/types/errors.go +++ b/x/pse/types/errors.go @@ -25,4 +25,7 @@ var ( // ErrOngoingDistribution is returned when a schedule update is attempted during an ongoing distribution. ErrOngoingDistribution = sdkerrors.Register(ModuleName, 8, "distribution is currently in progress") + + // ErrInvariantViolation is returned when an internal invariant is violated, disabling the module. + ErrInvariantViolation = sdkerrors.Register(ModuleName, 9, "invariant violation") ) From 44be1a74ba606c06951035b5b27c48983d09761c Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Thu, 12 Mar 2026 13:04:24 +0300 Subject: [PATCH 33/34] invariant: disable pse for non-positive community allocation --- x/pse/keeper/distribution.go | 24 ++++++++++-------------- x/pse/keeper/distribution_test.go | 26 +++++++------------------- 2 files changed, 17 insertions(+), 33 deletions(-) diff --git a/x/pse/keeper/distribution.go b/x/pse/keeper/distribution.go index 5d0b1361..9d3bc1ca 100644 --- a/x/pse/keeper/distribution.go +++ b/x/pse/keeper/distribution.go @@ -55,27 +55,23 @@ func (k Keeper) ProcessNextDistribution(ctx context.Context) error { sdkCtx := sdk.UnwrapSDKContext(ctx) - // If community allocation exists, start multi-block processing. + // Invariant: every distribution must have a positive community allocation + // (enforced at schedule creation, also checked here as safety net). communityAmount := getCommunityAllocationAmount(scheduledDistribution) - if communityAmount.IsPositive() { - if err := k.OngoingDistribution.Set(ctx, scheduledDistribution); err != nil { - return err - } - sdkCtx.Logger().Info("started multi-block community distribution", - "distribution_id", scheduledDistribution.ID, - "timestamp", scheduledDistribution.Timestamp) - return nil + if !communityAmount.IsPositive() { + return errorsmod.Wrapf( + types.ErrInvariantViolation, + "non-positive community allocation %s for distribution %d", + communityAmount, scheduledDistribution.ID, + ) } - // No community allocation — remove from schedule - if err := k.AllocationSchedule.Remove(ctx, scheduledDistribution.ID); err != nil { + if err := k.OngoingDistribution.Set(ctx, scheduledDistribution); err != nil { return err } - - sdkCtx.Logger().Info("processed single-block distribution", + sdkCtx.Logger().Info("started multi-block community distribution", "distribution_id", scheduledDistribution.ID, "timestamp", scheduledDistribution.Timestamp) - return nil } diff --git a/x/pse/keeper/distribution_test.go b/x/pse/keeper/distribution_test.go index e66ba4aa..6f200fdf 100644 --- a/x/pse/keeper/distribution_test.go +++ b/x/pse/keeper/distribution_test.go @@ -216,15 +216,15 @@ func TestDistribution_PrecisionWithMultipleRecipients(t *testing.T) { requireT.NoError(err) } - // Create and save distribution schedule - // Note: Community is excluded from this test since it has different distribution logic - // and is tested separately in other tests + // Create and save distribution schedule. + // Community allocation is required; non-community precision is the focus of this test. startTime := uint64(time.Now().Add(-1 * time.Hour).Unix()) schedule := []types.ScheduledDistribution{ { ID: 1, Timestamp: startTime, Allocations: []types.ClearingAccountAllocation{ + {ClearingAccount: types.ClearingAccountCommunity, Amount: allocationAmount}, {ClearingAccount: types.ClearingAccountFoundation, Amount: allocationAmount}, {ClearingAccount: types.ClearingAccountAlliance, Amount: allocationAmount}, {ClearingAccount: types.ClearingAccountPartnership, Amount: allocationAmount}, @@ -237,7 +237,7 @@ func TestDistribution_PrecisionWithMultipleRecipients(t *testing.T) { err = pseKeeper.SaveDistributionSchedule(ctx, schedule) requireT.NoError(err) - // Process distribution + // First call processes non-community allocations and starts multi-block community distribution. ctx = ctx.WithBlockTime(time.Unix(int64(startTime)+10, 0)) err = pseKeeper.ProcessNextDistribution(ctx) requireT.NoError(err) @@ -486,7 +486,7 @@ func TestDistribution_MultiBlockEndBlockerRouting(t *testing.T) { } // TestDistribution_NonCommunityOnlySingleBlock tests that a distribution with -// no community allocation completes in a single call to ProcessNextDistribution. +// zero community allocation triggers an invariant violation. func TestDistribution_NonCommunityOnlySingleBlock(t *testing.T) { requireT := require.New(t) @@ -543,21 +543,9 @@ func TestDistribution_NonCommunityOnlySingleBlock(t *testing.T) { }) requireT.NoError(err) - // Single call should complete everything + // Non-community allocations are processed, but zero community triggers invariant violation. err = pseKeeper.ProcessNextDistribution(ctx) - requireT.NoError(err) - - // No OngoingDistribution should be set (no community allocation) - _, err = pseKeeper.OngoingDistribution.Get(ctx) - requireT.ErrorIs(err, collections.ErrNotFound, "no OngoingDistribution for non-community-only distribution") - - // Schedule entry should be removed - _, err = pseKeeper.AllocationSchedule.Get(ctx, 1) - requireT.ErrorIs(err, collections.ErrNotFound, "schedule should be removed after single-block distribution") - - // Recipient should have received all non-community tokens - recipientBalance := bankKeeper.GetBalance(ctx, sdk.MustAccAddressFromBech32(recipientAddr), bondDenom) - requireT.Equal(amount.MulRaw(5).String(), recipientBalance.Amount.String()) + requireT.ErrorIs(err, types.ErrInvariantViolation) } // TestDistribution_EndBlockerWithScenarios mirrors TestKeeper_Distribute scenarios but routes From 3d8940e06f1241b4f33949a4d6a613a9999b0e87 Mon Sep 17 00:00:00 2001 From: metalarm10 Date: Thu, 12 Mar 2026 13:48:31 +0300 Subject: [PATCH 34/34] ci-fix: reconstruct TotalScore from account scores during InitGenesis --- x/pse/keeper/genesis.go | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/x/pse/keeper/genesis.go b/x/pse/keeper/genesis.go index 70b1b9fe..ff42e1a2 100644 --- a/x/pse/keeper/genesis.go +++ b/x/pse/keeper/genesis.go @@ -46,7 +46,8 @@ func (k Keeper) InitGenesis(ctx context.Context, genState types.GenesisState) er } } - // Populate account scores from genesis state. + // Populate account scores from genesis state and reconstruct TotalScore. + totalScores := make(map[uint64]sdkmath.Int) for _, accountScore := range genState.AccountScores { addr, err := k.addressCodec.StringToBytes(accountScore.Address) if err != nil { @@ -55,6 +56,16 @@ func (k Keeper) InitGenesis(ctx context.Context, genState types.GenesisState) er if err := k.SetDelegatorScore(ctx, accountScore.DistributionID, addr, accountScore.Score); err != nil { return err } + if existing, ok := totalScores[accountScore.DistributionID]; ok { + totalScores[accountScore.DistributionID] = existing.Add(accountScore.Score) + } else { + totalScores[accountScore.DistributionID] = accountScore.Score + } + } + for distID, total := range totalScores { + if err := k.TotalScore.Set(ctx, distID, total); err != nil { + return err + } } return k.DistributionDisabled.Set(ctx, genState.DistributionsDisabled)