Add voting power flag, file to treegen output

Author: jshufro

main.go

@@ -95,6 +95,12 @@ func main() {
 			Usage:   "Enable the rolling record capability of the Smartnode tree generator. Use this to store and load record caches instead of recalculating attestation performance each time you run treegen.",
 			Value:   false,
 		},
+		&cli.BoolFlag{
+			Name:    "generate-voting-power",
+			Aliases: []string{"gvp"},
+			Usage:   "If Enabled, a file containing the voting power breakdown of all nodes will be saved to the output directory.",
+			Value:   false,
+		},
 		&cli.StringFlag{
 			Name:    "cpuprofile",
 			Aliases: []string{"c"},

tree-gen.go

@@ -87,19 +87,20 @@ type treegenArguments struct {
 
 // Treegen holder for the requested execution metadata and necessary artifacts
 type treeGenerator struct {
-	log               *log.ColorLogger
-	errLog            *log.ColorLogger
-	rp                *rocketpool.RocketPool
-	cfg               *config.RocketPoolConfig
-	mgr               *state.NetworkStateManager
-	recordMgr         *rprewards.RollingRecordManager
-	bn                beacon.Client
-	beaconConfig      beacon.Eth2Config
-	targets           targets
-	outputDir         string
-	prettyPrint       bool
-	ruleset           uint64
-	useRollingRecords bool
+	log                 *log.ColorLogger
+	errLog              *log.ColorLogger
+	rp                  *rocketpool.RocketPool
+	cfg                 *config.RocketPoolConfig
+	mgr                 *state.NetworkStateManager
+	recordMgr           *rprewards.RollingRecordManager
+	bn                  beacon.Client
+	beaconConfig        beacon.Eth2Config
+	targets             targets
+	outputDir           string
+	prettyPrint         bool
+	ruleset             uint64
+	useRollingRecords   bool
+	generateVotingPower bool
 }
 
 // Generates a new rewards tree based on the command line flags
@@ -167,17 +168,18 @@ func GenerateTree(c *cli.Context) error {
 
 	// Create the generator
 	generator := treeGenerator{
-		log:               &logger,
-		errLog:            &errLogger,
-		rp:                rp,
-		cfg:               cfg,
-		bn:                bn,
-		mgr:               mgr,
-		beaconConfig:      beaconConfig,
-		outputDir:         c.String("output-dir"),
-		prettyPrint:       c.Bool("pretty-print"),
-		ruleset:           c.Uint64("ruleset"),
-		useRollingRecords: c.Bool("use-rolling-records"),
+		log:                 &logger,
+		errLog:              &errLogger,
+		rp:                  rp,
+		cfg:                 cfg,
+		bn:                  bn,
+		mgr:                 mgr,
+		beaconConfig:        beaconConfig,
+		outputDir:           c.String("output-dir"),
+		prettyPrint:         c.Bool("pretty-print"),
+		ruleset:             c.Uint64("ruleset"),
+		useRollingRecords:   c.Bool("use-rolling-records"),
+		generateVotingPower: c.Bool("generate-voting-power"),
 	}
 
 	// initialize the generator targets
@@ -414,6 +416,14 @@ func (g *treeGenerator) generateRewardsFile(treegen *rprewards.TreeGenerator) (r
 	return treegen.GenerateTreeWithRuleset(g.ruleset)
 }
 
+func (g *treeGenerator) serializeVotingPower(votingPowerFile *VotingPowerFile) ([]byte, error) {
+	if g.prettyPrint {
+		return json.MarshalIndent(votingPowerFile, "", "\t")
+	}
+
+	return json.Marshal(votingPowerFile)
+}
+
 // Serializes the minipool performance file into JSON
 func (g *treeGenerator) serializeMinipoolPerformance(rewardsFile rprewards.IRewardsFile) ([]byte, error) {
 	perfFile := rewardsFile.GetMinipoolPerformanceFile()
@@ -435,7 +445,7 @@ func (g *treeGenerator) serializeRewardsTree(rewardsFile rprewards.IRewardsFile)
 }
 
 // Writes both the performance file and the rewards file to disk
-func (g *treeGenerator) writeFiles(rewardsFile rprewards.IRewardsFile) error {
+func (g *treeGenerator) writeFiles(rewardsFile rprewards.IRewardsFile, votingPowerFile *VotingPowerFile) error {
 	g.log.Printlnf("Saving JSON files...")
 	index := rewardsFile.GetHeader().Index
 
@@ -472,6 +482,21 @@ func (g *treeGenerator) writeFiles(rewardsFile rprewards.IRewardsFile) error {
 	}
 
 	g.log.Printlnf("Saved rewards snapshot file to %s", rewardsTreePath)
+
+	// Write the voting power file to disk
+	if votingPowerFile != nil {
+		votingPowerFileBytes, err := g.serializeVotingPower(votingPowerFile)
+		if err != nil {
+			return fmt.Errorf("error serializing voting power file into JSON: %w", err)
+		}
+
+		votingFilePath := filepath.Join(g.outputDir, fmt.Sprintf("rp-voting-power-%s-%d.json", string(g.cfg.Smartnode.Network.Value.(cfgtypes.Network)), votingPowerFile.ConsensusSlot))
+		err = os.WriteFile(votingFilePath, votingPowerFileBytes, 0644)
+		if err != nil {
+			return fmt.Errorf("error saving voting power file to %s: %w", votingFilePath, err)
+		}
+		g.log.Printlnf("Saved voting power file to %s", votingFilePath)
+	}
 	g.log.Printlnf("Successfully generated rewards snapshot for interval %d", index)
 
 	return nil
@@ -622,6 +647,7 @@ func (g *treeGenerator) approximateRethSpRewards() error {
 
 // Generate a complete rewards tree
 func (g *treeGenerator) generateTree() error {
+	var votingPowerFile *VotingPowerFile
 	args, err := g.getTreegenArgs()
 	if err != nil {
 		return fmt.Errorf("error compiling treegen arguments: %w", err)
@@ -633,6 +659,12 @@ func (g *treeGenerator) generateTree() error {
 		return err
 	}
 
+	// If a voting power file was requested, generate it now.
+	if g.generateVotingPower {
+		votingPowerFile = g.GenerateVotingPower(args.state)
+		votingPowerFile.Time = args.endTime
+	}
+
 	// Generate the rewards file
 	start := time.Now()
 	rewardsFile, err := g.generateRewardsFile(treegen)
@@ -656,7 +688,7 @@ func (g *treeGenerator) generateTree() error {
 		}
 	}
 
-	err = g.writeFiles(rewardsFile)
+	err = g.writeFiles(rewardsFile, votingPowerFile)
 	if err != nil {
 		return err
 	}

voting-power.go

@@ -0,0 +1,92 @@
+package main
+
+import (
+	"math/big"
+	"time"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/rocket-pool/smartnode/shared/services/rewards"
+	"github.com/rocket-pool/smartnode/shared/services/state"
+	cfgtypes "github.com/rocket-pool/smartnode/shared/types/config"
+)
+
+var (
+	oneEth   = big.NewInt(1e18)
+	_1_5_Eth = big.NewInt(15e17)
+)
+
+type VotingPowerFile struct {
+	Network        string                                   `json:"network"`
+	Time           time.Time                                `json:"time"`
+	ConsensusSlot  uint64                                   `json:"consensusSlot"`
+	ExecutionBlock uint64                                   `json:"executionBlock"`
+	TotalPower     *rewards.QuotedBigInt                    `json:"totalPower"`
+	NodePower      map[common.Address]*rewards.QuotedBigInt `json:"nodePower"`
+}
+
+func getNodeVotingPower(s *state.NetworkState, ethProvided *big.Int, nodeStake *big.Int) *big.Int {
+	rplPrice := s.NetworkDetails.RplPrice
+
+	// No RPL staked means no voting power
+	if nodeStake.Sign() == 0 {
+		return big.NewInt(0)
+	}
+
+	// First calculate the maximum rpl that can be used as input
+	// maxVotingRpl := (eligibleBondedEth * 1.5 Eth / RplPrice)
+	maxVotingRpl := big.NewInt(0)
+	maxVotingRpl.Mul(ethProvided, _1_5_Eth)
+	maxVotingRpl.Quo(maxVotingRpl, rplPrice)
+
+	// Determine the voting RPL
+	// votingRpl := min(maxVotingRpl, nodeStake)
+	var votingRpl *big.Int
+	if maxVotingRpl.Cmp(nodeStake) <= 0 {
+		votingRpl = maxVotingRpl
+	} else {
+		votingRpl = nodeStake
+	}
+
+	// Now take the square root and divide by 2
+	// Because the units are in wei, we need to multiply votingRpl by 1 Eth before square rooting.
+	votingPower := big.NewInt(0)
+	votingPower.Mul(votingRpl, oneEth)
+	votingPower.Sqrt(votingPower)
+	votingPower.Rsh(votingPower, 1)
+	return votingPower
+
+}
+
+func (g *treeGenerator) GenerateVotingPower(s *state.NetworkState) *VotingPowerFile {
+	out := new(VotingPowerFile)
+
+	out.Network = string(g.cfg.Smartnode.Network.Value.(cfgtypes.Network))
+	out.ConsensusSlot = s.BeaconSlotNumber
+	out.ExecutionBlock = s.ElBlockNumber
+	out.TotalPower = rewards.NewQuotedBigInt(0)
+	out.NodePower = make(map[common.Address]*rewards.QuotedBigInt, len(s.NodeDetails))
+	for _, node := range s.NodeDetails {
+		activeMinipoolCount := int64(0)
+		for _, mpd := range s.MinipoolDetailsByNode[node.NodeAddress] {
+			// Ignore finalised
+			if mpd.Finalised {
+				continue
+			}
+
+			activeMinipoolCount += 1
+		}
+
+		// Get provided ETH (32 * minipoolCount - matched)
+		ethProvided := big.NewInt(activeMinipoolCount * 32)
+		ethProvided.Mul(ethProvided, oneEth)
+		ethProvided.Sub(ethProvided, node.EthMatched)
+
+		// Calculate the Voting Power
+		nodeVotingPower := rewards.NewQuotedBigInt(0)
+		nodeVotingPower.Set(getNodeVotingPower(s, ethProvided, node.RplStake))
+		out.TotalPower.Add(&out.TotalPower.Int, &nodeVotingPower.Int)
+		out.NodePower[node.NodeAddress] = nodeVotingPower
+	}
+
+	return out
+}