[OTE-695] Aggregate referred volume in endblocker (#2218)

dydxprotocol · Sep 16, 2024 · d964f26 · d964f26
1 parent 1c7884e
commit d964f26
Show file tree

Hide file tree

Showing 5 changed files with 180 additions and 0 deletions.
diff --git a/protocol/x/affiliates/abci.go b/protocol/x/affiliates/abci.go
@@ -0,0 +1,19 @@
+package affiliates
+
+import (
+	"runtime/debug"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/dydxprotocol/v4-chain/protocol/lib/log"
+	"github.com/dydxprotocol/v4-chain/protocol/x/affiliates/keeper"
+)
+
+func EndBlocker(
+	ctx sdk.Context,
+	keeper *keeper.Keeper,
+) {
+	if err := keeper.AggregateAffiliateReferredVolumeForFills(ctx); err != nil {
+		log.ErrorLog(ctx, "error aggregating affiliate volume for fills", "error",
+			err, "stack", string(debug.Stack()))
+	}
+}
diff --git a/protocol/x/affiliates/keeper/keeper.go b/protocol/x/affiliates/keeper/keeper.go
@@ -258,3 +258,36 @@ func (k *Keeper) SetRevShareKeeper(revShareKeeper types.RevShareKeeper) {
 func (k Keeper) GetIndexerEventManager() indexer_manager.IndexerEventManager {
 	return k.indexerEventManager
 }
+
+func (k Keeper) AggregateAffiliateReferredVolumeForFills(
+	ctx sdk.Context,
+) error {
+	blockStats := k.statsKeeper.GetBlockStats(ctx)
+	referredByCache := make(map[string]string)
+	for _, fill := range blockStats.Fills {
+		// Add taker's referred volume to the cache
+		if _, ok := referredByCache[fill.Taker]; !ok {
+			referredByAddrTaker, found := k.GetReferredBy(ctx, fill.Taker)
+			if !found {
+				continue
+			}
+			referredByCache[fill.Taker] = referredByAddrTaker
+		}
+		if err := k.AddReferredVolume(ctx, referredByCache[fill.Taker], lib.BigU(fill.Notional)); err != nil {
+			return err
+		}
+
+		// Add maker's referred volume to the cache
+		if _, ok := referredByCache[fill.Maker]; !ok {
+			referredByAddrMaker, found := k.GetReferredBy(ctx, fill.Maker)
+			if !found {
+				continue
+			}
+			referredByCache[fill.Maker] = referredByAddrMaker
+		}
+		if err := k.AddReferredVolume(ctx, referredByCache[fill.Maker], lib.BigU(fill.Notional)); err != nil {
+			return err
+		}
+	}
+	return nil
+}
diff --git a/protocol/x/affiliates/keeper/keeper_test.go b/protocol/x/affiliates/keeper/keeper_test.go
@@ -16,6 +16,7 @@ import (
 	keepertest "github.com/dydxprotocol/v4-chain/protocol/testutil/keeper"
 	"github.com/dydxprotocol/v4-chain/protocol/x/affiliates/keeper"
 	"github.com/dydxprotocol/v4-chain/protocol/x/affiliates/types"
+	statskeeper "github.com/dydxprotocol/v4-chain/protocol/x/stats/keeper"
 	statstypes "github.com/dydxprotocol/v4-chain/protocol/x/stats/types"
 	"github.com/stretchr/testify/require"
 )
@@ -392,3 +393,120 @@ func TestRegisterAffiliateEmitEvent(t *testing.T) {
 	require.Equal(t, 1, len(events))
 	require.Equal(t, expectedEvent, events[0])
 }
+
+func TestAggregateAffiliateReferredVolumeForFills(t *testing.T) {
+	affiliate := constants.AliceAccAddress.String()
+	referee1 := constants.BobAccAddress.String()
+	referee2 := constants.DaveAccAddress.String()
+	maker := constants.CarlAccAddress.String()
+	testCases := []struct {
+		name           string
+		referrals      int
+		expectedVolume *big.Int
+		setup          func(t *testing.T, ctx sdk.Context, k *keeper.Keeper, statsKeeper *statskeeper.Keeper)
+	}{
+		{
+			name:           "0 referrals",
+			expectedVolume: big.NewInt(0),
+			setup: func(t *testing.T, ctx sdk.Context, k *keeper.Keeper, statsKeeper *statskeeper.Keeper) {
+				statsKeeper.SetBlockStats(ctx, &statstypes.BlockStats{
+					Fills: []*statstypes.BlockStats_Fill{
+						{
+							Taker:    referee1,
+							Maker:    maker,
+							Notional: 100_000_000_000,
+						},
+					},
+				})
+			},
+		},
+		{
+			name:           "1 referral",
+			referrals:      1,
+			expectedVolume: big.NewInt(100_000_000_000),
+			setup: func(t *testing.T, ctx sdk.Context, k *keeper.Keeper, statsKeeper *statskeeper.Keeper) {
+				err := k.RegisterAffiliate(ctx, referee1, affiliate)
+				require.NoError(t, err)
+				statsKeeper.SetBlockStats(ctx, &statstypes.BlockStats{
+					Fills: []*statstypes.BlockStats_Fill{
+						{
+							Taker:    referee1,
+							Maker:    maker,
+							Notional: 100_000_000_000,
+						},
+					},
+				})
+			},
+		},
+		{
+			name:           "2 referrals",
+			referrals:      2,
+			expectedVolume: big.NewInt(300_000_000_000),
+			setup: func(t *testing.T, ctx sdk.Context, k *keeper.Keeper, statsKeeper *statskeeper.Keeper) {
+				err := k.RegisterAffiliate(ctx, referee1, affiliate)
+				require.NoError(t, err)
+				err = k.RegisterAffiliate(ctx, referee2, affiliate)
+				require.NoError(t, err)
+				statsKeeper.SetBlockStats(ctx, &statstypes.BlockStats{
+					Fills: []*statstypes.BlockStats_Fill{
+						{
+							Taker:    referee1,
+							Maker:    maker,
+							Notional: 100_000_000_000,
+						},
+						{
+							Taker:    referee2,
+							Maker:    maker,
+							Notional: 200_000_000_000,
+						},
+					},
+				})
+			},
+		},
+		{
+			name:           "2 referrals, maker also referred",
+			referrals:      2,
+			expectedVolume: big.NewInt(600_000_000_000),
+			setup: func(t *testing.T, ctx sdk.Context, k *keeper.Keeper, statsKeeper *statskeeper.Keeper) {
+				err := k.RegisterAffiliate(ctx, referee1, affiliate)
+				require.NoError(t, err)
+				err = k.RegisterAffiliate(ctx, referee2, affiliate)
+				require.NoError(t, err)
+				err = k.RegisterAffiliate(ctx, maker, affiliate)
+				require.NoError(t, err)
+				statsKeeper.SetBlockStats(ctx, &statstypes.BlockStats{
+					Fills: []*statstypes.BlockStats_Fill{
+						{
+							Taker:    referee1,
+							Maker:    maker,
+							Notional: 100_000_000_000,
+						},
+						{
+							Taker:    referee2,
+							Maker:    maker,
+							Notional: 200_000_000_000,
+						},
+					},
+				})
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			tApp := testapp.NewTestAppBuilder(t).Build()
+			ctx := tApp.InitChain()
+			k := tApp.App.AffiliatesKeeper
+			statsKeeper := tApp.App.StatsKeeper
+
+			tc.setup(t, ctx, &k, &statsKeeper)
+
+			err := k.AggregateAffiliateReferredVolumeForFills(ctx)
+			require.NoError(t, err)
+
+			referredVolume, err := k.GetReferredVolume(ctx, affiliate)
+			require.NoError(t, err)
+			require.Equal(t, tc.expectedVolume, referredVolume)
+		})
+	}
+}
diff --git a/protocol/x/affiliates/module.go b/protocol/x/affiliates/module.go
@@ -14,6 +14,7 @@ import (
 	cdctypes "github.com/cosmos/cosmos-sdk/codec/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
+	"github.com/dydxprotocol/v4-chain/protocol/lib"
 	"github.com/dydxprotocol/v4-chain/protocol/x/affiliates/client/cli"
 	"github.com/dydxprotocol/v4-chain/protocol/x/affiliates/keeper"
 	"github.com/dydxprotocol/v4-chain/protocol/x/affiliates/types"
@@ -141,3 +142,10 @@ func (am AppModule) ExportGenesis(ctx sdk.Context, cdc codec.JSONCodec) json.Raw
 // consensus-breaking change introduced by the module. To avoid wrong/empty versions, the initial version should
 // be set to 1.
 func (AppModule) ConsensusVersion() uint64 { return 1 }
+
+func (am AppModule) EndBlock(ctx sdk.Context) {
+	EndBlocker(
+		lib.UnwrapSDKContext(ctx, types.ModuleName),
+		&am.keeper,
+	)
+}
diff --git a/protocol/x/affiliates/types/expected_keepers.go b/protocol/x/affiliates/types/expected_keepers.go
@@ -5,10 +5,12 @@ import (
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	revsharetypes "github.com/dydxprotocol/v4-chain/protocol/x/revshare/types"
+	stattypes "github.com/dydxprotocol/v4-chain/protocol/x/stats/types"
 )
 
 type StatsKeeper interface {
 	GetStakedAmount(ctx sdk.Context, delegatorAddr string) *big.Int
+	GetBlockStats(ctx sdk.Context) *stattypes.BlockStats
 }
 
 type RevShareKeeper interface {