@@ -87,19 +87,20 @@ type treegenArguments struct {
8787
8888// Treegen holder for the requested execution metadata and necessary artifacts
8989type treeGenerator struct {
90- log * log.ColorLogger
91- errLog * log.ColorLogger
92- rp * rocketpool.RocketPool
93- cfg * config.RocketPoolConfig
94- mgr * state.NetworkStateManager
95- recordMgr * rprewards.RollingRecordManager
96- bn beacon.Client
97- beaconConfig beacon.Eth2Config
98- targets targets
99- outputDir string
100- prettyPrint bool
101- ruleset uint64
102- useRollingRecords bool
90+ log * log.ColorLogger
91+ errLog * log.ColorLogger
92+ rp * rocketpool.RocketPool
93+ cfg * config.RocketPoolConfig
94+ mgr * state.NetworkStateManager
95+ recordMgr * rprewards.RollingRecordManager
96+ bn beacon.Client
97+ beaconConfig beacon.Eth2Config
98+ targets targets
99+ outputDir string
100+ prettyPrint bool
101+ ruleset uint64
102+ useRollingRecords bool
103+ generateVotingPower bool
103104}
104105
105106// Generates a new rewards tree based on the command line flags
@@ -170,17 +171,18 @@ func GenerateTree(c *cli.Context) error {
170171
171172 // Create the generator
172173 generator := treeGenerator {
173- log : & logger ,
174- errLog : & errLogger ,
175- rp : rp ,
176- cfg : cfg ,
177- bn : bn ,
178- mgr : mgr ,
179- beaconConfig : beaconConfig ,
180- outputDir : c .String ("output-dir" ),
181- prettyPrint : c .Bool ("pretty-print" ),
182- ruleset : c .Uint64 ("ruleset" ),
183- useRollingRecords : c .Bool ("use-rolling-records" ),
174+ log : & logger ,
175+ errLog : & errLogger ,
176+ rp : rp ,
177+ cfg : cfg ,
178+ bn : bn ,
179+ mgr : mgr ,
180+ beaconConfig : beaconConfig ,
181+ outputDir : c .String ("output-dir" ),
182+ prettyPrint : c .Bool ("pretty-print" ),
183+ ruleset : c .Uint64 ("ruleset" ),
184+ useRollingRecords : c .Bool ("use-rolling-records" ),
185+ generateVotingPower : c .Bool ("generate-voting-power" ),
184186 }
185187
186188 // initialize the generator targets
@@ -417,6 +419,14 @@ func (g *treeGenerator) generateRewardsFile(treegen *rprewards.TreeGenerator) (r
417419 return treegen .GenerateTreeWithRuleset (g .ruleset )
418420}
419421
422+ func (g * treeGenerator ) serializeVotingPower (votingPowerFile * VotingPowerFile ) ([]byte , error ) {
423+ if g .prettyPrint {
424+ return json .MarshalIndent (votingPowerFile , "" , "\t " )
425+ }
426+
427+ return json .Marshal (votingPowerFile )
428+ }
429+
420430// Serializes the minipool performance file into JSON
421431func (g * treeGenerator ) serializeMinipoolPerformance (rewardsFile rprewards.IRewardsFile ) ([]byte , error ) {
422432 perfFile := rewardsFile .GetMinipoolPerformanceFile ()
@@ -438,7 +448,7 @@ func (g *treeGenerator) serializeRewardsTree(rewardsFile rprewards.IRewardsFile)
438448}
439449
440450// Writes both the performance file and the rewards file to disk
441- func (g * treeGenerator ) writeFiles (rewardsFile rprewards.IRewardsFile ) error {
451+ func (g * treeGenerator ) writeFiles (rewardsFile rprewards.IRewardsFile , votingPowerFile * VotingPowerFile ) error {
442452 g .log .Printlnf ("Saving JSON files..." )
443453 index := rewardsFile .GetHeader ().Index
444454
@@ -475,6 +485,21 @@ func (g *treeGenerator) writeFiles(rewardsFile rprewards.IRewardsFile) error {
475485 }
476486
477487 g .log .Printlnf ("Saved rewards snapshot file to %s" , rewardsTreePath )
488+
489+ // Write the voting power file to disk
490+ if votingPowerFile != nil {
491+ votingPowerFileBytes , err := g .serializeVotingPower (votingPowerFile )
492+ if err != nil {
493+ return fmt .Errorf ("error serializing voting power file into JSON: %w" , err )
494+ }
495+
496+ votingFilePath := filepath .Join (g .outputDir , fmt .Sprintf ("rp-voting-power-%s-%d.json" , string (g .cfg .Smartnode .Network .Value .(cfgtypes.Network )), votingPowerFile .ConsensusSlot ))
497+ err = os .WriteFile (votingFilePath , votingPowerFileBytes , 0644 )
498+ if err != nil {
499+ return fmt .Errorf ("error saving voting power file to %s: %w" , votingFilePath , err )
500+ }
501+ g .log .Printlnf ("Saved voting power file to %s" , votingFilePath )
502+ }
478503 g .log .Printlnf ("Successfully generated rewards snapshot for interval %d" , index )
479504
480505 return nil
@@ -625,6 +650,7 @@ func (g *treeGenerator) approximateRethSpRewards() error {
625650
626651// Generate a complete rewards tree
627652func (g * treeGenerator ) generateTree () error {
653+ var votingPowerFile * VotingPowerFile
628654 args , err := g .getTreegenArgs ()
629655 if err != nil {
630656 return fmt .Errorf ("error compiling treegen arguments: %w" , err )
@@ -636,6 +662,12 @@ func (g *treeGenerator) generateTree() error {
636662 return err
637663 }
638664
665+ // If a voting power file was requested, generate it now.
666+ if g .generateVotingPower {
667+ votingPowerFile = g .GenerateVotingPower (args .state )
668+ votingPowerFile .Time = args .endTime
669+ }
670+
639671 // Generate the rewards file
640672 start := time .Now ()
641673 rewardsFile , err := g .generateRewardsFile (treegen )
@@ -659,7 +691,7 @@ func (g *treeGenerator) generateTree() error {
659691 }
660692 }
661693
662- err = g .writeFiles (rewardsFile )
694+ err = g .writeFiles (rewardsFile , votingPowerFile )
663695 if err != nil {
664696 return err
665697 }
0 commit comments