go-cyber/x/bandwidth/keeper/keeper.go

package keeper

import (
	sdkmath "cosmossdk.io/math"
	"fmt"
	storetypes "github.com/cosmos/cosmos-sdk/store/types"
	ctypes "github.com/cybercongress/go-cyber/v7/types"

	"github.com/cometbft/cometbft/libs/log"
	"github.com/cosmos/cosmos-sdk/codec"
	"github.com/cosmos/cosmos-sdk/telemetry"
	sdk "github.com/cosmos/cosmos-sdk/types"

	"github.com/cybercongress/go-cyber/v7/x/bandwidth/types"
	gtypes "github.com/cybercongress/go-cyber/v7/x/graph/types"
)

type BandwidthMeter struct {
	stakeProvider types.AccountStakeProvider
	cdc           codec.BinaryCodec
	storeKey      storetypes.StoreKey
	tkey          *storetypes.TransientStoreKey

	currentCreditPrice         sdk.Dec
	bandwidthSpentByBlock      map[uint64]uint64
	totalSpentForSlidingWindow uint64

	authority string
}

func NewBandwidthMeter(
	cdc codec.BinaryCodec,
	key storetypes.StoreKey,
	tkey *storetypes.TransientStoreKey,
	asp types.AccountStakeProvider,
	authority string,
) *BandwidthMeter {
	return &BandwidthMeter{
		cdc:                   cdc,
		storeKey:              key,
		tkey:                  tkey,
		stakeProvider:         asp,
		bandwidthSpentByBlock: make(map[uint64]uint64),
		authority:             authority,
	}
}

func (bm BandwidthMeter) GetAuthority() string { return bm.authority }

func (bm BandwidthMeter) Logger(ctx sdk.Context) log.Logger {
	return ctx.Logger().With("module", fmt.Sprintf("x/%s", types.ModuleName))
}

func (bm BandwidthMeter) SetParams(ctx sdk.Context, p types.Params) error {
	if err := p.Validate(); err != nil {
		return err
	}

	store := ctx.KVStore(bm.storeKey)
	bz := bm.cdc.MustMarshal(&p)
	store.Set(types.ParamsKey, bz)

	return nil
}

func (bm BandwidthMeter) GetParams(ctx sdk.Context) (p types.Params) {
	store := ctx.KVStore(bm.storeKey)
	bz := store.Get(types.ParamsKey)
	if bz == nil {
		return p
	}

	bm.cdc.MustUnmarshal(bz, &p)
	return p
}

func (bm *BandwidthMeter) LoadState(ctx sdk.Context) {
	params := bm.GetParams(ctx)
	bm.totalSpentForSlidingWindow = 0
	// TODO test case when period parameter is increased
	bm.bandwidthSpentByBlock = bm.GetValuesForPeriod(ctx, params.RecoveryPeriod)
	for _, spentBandwidth := range bm.bandwidthSpentByBlock {
		bm.totalSpentForSlidingWindow += spentBandwidth
	}
	bm.currentCreditPrice = bm.GetBandwidthPrice(ctx, params.BasePrice)
}

func (bm *BandwidthMeter) InitState() {
	bm.totalSpentForSlidingWindow = 0

	window := make(map[uint64]uint64)
	window[1] = 0
	bm.bandwidthSpentByBlock = window
}

func (bm BandwidthMeter) GetBandwidthPrice(ctx sdk.Context, basePrice sdk.Dec) sdk.Dec {
	store := ctx.KVStore(bm.storeKey)
	priceAsBytes := store.Get(types.LastBandwidthPrice)
	if priceAsBytes == nil {
		return basePrice
	}
	var price types.Price
	bm.cdc.MustUnmarshal(priceAsBytes, &price)
	return price.Price
}

func (bm BandwidthMeter) StoreBandwidthPrice(ctx sdk.Context, price sdk.Dec) {
	store := ctx.KVStore(bm.storeKey)
	store.Set(types.LastBandwidthPrice, bm.cdc.MustMarshal(&types.Price{Price: price}))
}

func (bm BandwidthMeter) GetDesirableBandwidth(ctx sdk.Context) uint64 {
	store := ctx.KVStore(bm.storeKey)
	bandwidthAsBytes := store.Get(types.TotalBandwidth)
	if bandwidthAsBytes == nil {
		return 0
	}
	return sdk.BigEndianToUint64(bandwidthAsBytes)
}

func (bm BandwidthMeter) AddToDesirableBandwidth(ctx sdk.Context, toAdd uint64) {
	current := bm.GetDesirableBandwidth(ctx)
	store := ctx.KVStore(bm.storeKey)
	store.Set(types.TotalBandwidth, sdk.Uint64ToBigEndian(current+toAdd))
}

func (bm BandwidthMeter) SetDesirableBandwidth(ctx sdk.Context, value uint64) {
	store := ctx.KVStore(bm.storeKey)
	store.Set(types.TotalBandwidth, sdk.Uint64ToBigEndian(value))
}

func (bm *BandwidthMeter) AddToBlockBandwidth(ctx sdk.Context, value uint64) {
	store := ctx.TransientStore(bm.tkey)

	current := uint64(0)
	blockBandwidth := store.Get(types.BlockBandwidth)
	if blockBandwidth != nil {
		current = sdk.BigEndianToUint64(blockBandwidth)
	}

	store.Set(types.BlockBandwidth, sdk.Uint64ToBigEndian(current+value))
}

// Here we move bandwidth window:
// Remove first block of window and add new block to window end
func (bm *BandwidthMeter) CommitBlockBandwidth(ctx sdk.Context) {
	params := bm.GetParams(ctx)

	tStore := ctx.TransientStore(bm.tkey)
	currentBlockSpentBandwidth := sdk.BigEndianToUint64(tStore.Get(types.BlockBandwidth))

	defer func() {
		if currentBlockSpentBandwidth > 0 {
			bm.Logger(ctx).Info("Block", "bandwidth", currentBlockSpentBandwidth)
			bm.Logger(ctx).Info("Window", "bandwidth", bm.totalSpentForSlidingWindow)
		}

		telemetry.SetGauge(float32(currentBlockSpentBandwidth), types.ModuleName, "block_bandwidth")
		telemetry.SetGauge(float32(bm.totalSpentForSlidingWindow), types.ModuleName, "window_bandwidth")
		tStore.Set(types.BlockBandwidth, sdk.Uint64ToBigEndian(0))
	}()

	bm.totalSpentForSlidingWindow += currentBlockSpentBandwidth

	newWindowEnd := ctx.BlockHeight()
	windowStart := newWindowEnd - int64(params.RecoveryPeriod)
	if windowStart < 0 {
		windowStart = 0
	}

	// clean window slot in in-memory
	windowStartValue, exists := bm.bandwidthSpentByBlock[uint64(windowStart)]
	if exists {
		bm.totalSpentForSlidingWindow -= windowStartValue
		delete(bm.bandwidthSpentByBlock, uint64(windowStart))
	}

	// clean window slot in storage
	// TODO will be removed, need to index blocks bandwidth
	store := ctx.KVStore(bm.storeKey)
	if store.Has(types.BlockStoreKey(uint64(windowStart))) {
		store.Delete(types.BlockStoreKey(uint64(windowStart)))
	}

	bm.SetBlockBandwidth(ctx, uint64(ctx.BlockHeight()), currentBlockSpentBandwidth)
	bm.bandwidthSpentByBlock[uint64(newWindowEnd)] = currentBlockSpentBandwidth
}

func (bm *BandwidthMeter) GetCurrentBlockSpentBandwidth(ctx sdk.Context) uint64 {
	tstore := ctx.TransientStore(bm.tkey)
	return sdk.BigEndianToUint64(tstore.Get(types.BlockBandwidth))
}

func (bm *BandwidthMeter) GetCurrentNetworkLoad(ctx sdk.Context) sdk.Dec {
	return sdk.NewDec(int64(bm.totalSpentForSlidingWindow)).QuoInt64(int64(bm.GetDesirableBandwidth(ctx)))
}

func (bm *BandwidthMeter) GetMaxBlockBandwidth(ctx sdk.Context) uint64 {
	params := bm.GetParams(ctx)
	maxBlockBandwidth := params.MaxBlockBandwidth
	return maxBlockBandwidth
}

func (bm *BandwidthMeter) GetCurrentCreditPrice() sdk.Dec {
	return bm.currentCreditPrice
}

// Adjusting of bandwidth price is disabled after the upgrade to v6, set to base price and never changes
func (bm *BandwidthMeter) AdjustPrice(ctx sdk.Context) {
	params := bm.GetParams(ctx)

	desirableBandwidth := bm.GetDesirableBandwidth(ctx)
	baseBandwidth := params.BaseLoad.MulInt64(int64(desirableBandwidth)).RoundInt64()
	if baseBandwidth != 0 {
		telemetry.SetGauge(float32(bm.totalSpentForSlidingWindow)/float32(baseBandwidth), types.ModuleName, "load")

		load := sdk.NewDec(int64(bm.totalSpentForSlidingWindow)).QuoInt64(baseBandwidth)
		bm.Logger(ctx).Info("Load", "value", load.String())

		newPrice := params.BasePrice
		bm.Logger(ctx).Info("Price", "value", newPrice.String())
		telemetry.SetGauge(float32(newPrice.MulInt64(1000).RoundInt64()), types.ModuleName, "price")

		bm.currentCreditPrice = newPrice
		bm.StoreBandwidthPrice(ctx, newPrice)
	}
}

func (bm *BandwidthMeter) GetTotalCyberlinksCost(ctx sdk.Context, tx sdk.Tx) uint64 {
	bandwidthForTx := uint64(0)
	for _, msg := range tx.GetMsgs() {
		linkMsg := msg.(*gtypes.MsgCyberlink)
		bandwidthForTx += uint64(len(linkMsg.Links)) * 1000
	}
	return bandwidthForTx
}

func (bm *BandwidthMeter) BurnAccountBandwidthVolt(ctx sdk.Context, amt uint64, address sdk.AccAddress) error {
	coin := sdk.NewCoin(ctypes.VOLT, sdkmath.NewIntFromUint64(amt))
	if err := bm.stakeProvider.SendCoinsFromAccountToModule(ctx, address, types.BandwidthName, sdk.NewCoins(coin)); err != nil {
		return err
	}
	if err := bm.stakeProvider.BurnCoins(ctx, types.BandwidthName, sdk.NewCoins(coin)); err != nil {
		return err
	}
	return nil
}

func (bm *BandwidthMeter) GetCurrentVoltsAccountBandwidth(ctx sdk.Context, address sdk.AccAddress) types.NeuronBandwidth {
	balance := bm.stakeProvider.GetAccountStakeVolt(ctx, address)
	return types.NeuronBandwidth{
		Neuron:           address.String(),
		RemainedValue:    uint64(balance),
		MaxValue:         uint64(balance),
		LastUpdatedBlock: uint64(ctx.BlockHeight()),
	}
}

func (bm *BandwidthMeter) SetZeroAccountBandwidth(ctx sdk.Context, address sdk.AccAddress) {
	bw := types.NeuronBandwidth{
		Neuron:           address.String(),
		RemainedValue:    0,
		MaxValue:         0,
		LastUpdatedBlock: uint64(ctx.BlockHeight()),
	}

	ctx.KVStore(bm.storeKey).Set(types.AccountStoreKey(bw.Neuron), bm.cdc.MustMarshal(&bw))
}

func (bm *BandwidthMeter) HasEnoughAccountBandwidthVolt(ctx sdk.Context, toConsume uint64, address sdk.AccAddress) bool {
	return uint64(bm.stakeProvider.GetAccountStakeVolt(ctx, address)) >= toConsume
}