@@ -88,19 +88,19 @@ type treegenArguments struct {
8888
8989// Treegen holder for the requested execution metadata and necessary artifacts
9090type treeGenerator struct {
91- log * log.ColorLogger
92- errLog * log.ColorLogger
93- rp rprewards.RewardsExecutionClient
94- rpNative * rocketpool.RocketPool
95- cfg * config.RocketPoolConfig
96- mgr * state.NetworkStateManager
97- bn beacon.Client
98- beaconConfig beacon.Eth2Config
99- targets targets
100- outputDir string
101- prettyPrint bool
102- ruleset uint64
103- useRollingRecords bool
91+ log * log.ColorLogger
92+ errLog * log.ColorLogger
93+ rp rprewards.RewardsExecutionClient
94+ rpNative * rocketpool.RocketPool
95+ cfg * config.RocketPoolConfig
96+ mgr * state.NetworkStateManager
97+ bn beacon.Client
98+ beaconConfig beacon.Eth2Config
99+ targets targets
100+ outputDir string
101+ prettyPrint bool
102+ ruleset uint64
103+ generateVotingPower bool
104104}
105105
106106// Generates a new rewards tree based on the command line flags
@@ -168,18 +168,18 @@ func GenerateTree(c *cli.Context) error {
168168
169169 // Create the generator
170170 generator := treeGenerator {
171- log : & logger ,
172- errLog : & errLogger ,
173- rp : rprewards .NewRewardsExecutionClient (rp ),
174- rpNative : rp ,
175- cfg : cfg ,
176- bn : bn ,
177- mgr : mgr ,
178- beaconConfig : beaconConfig ,
179- outputDir : c .String ("output-dir" ),
180- prettyPrint : c .Bool ("pretty-print" ),
181- ruleset : c .Uint64 ("ruleset" ),
182- useRollingRecords : c .Bool ("use-rolling-records " ),
171+ log : & logger ,
172+ errLog : & errLogger ,
173+ rp : rprewards .NewRewardsExecutionClient (rp ),
174+ rpNative : rp ,
175+ cfg : cfg ,
176+ bn : bn ,
177+ mgr : mgr ,
178+ beaconConfig : beaconConfig ,
179+ outputDir : c .String ("output-dir" ),
180+ prettyPrint : c .Bool ("pretty-print" ),
181+ ruleset : c .Uint64 ("ruleset" ),
182+ generateVotingPower : c .Bool ("generate-voting-power " ),
183183 }
184184
185185 // initialize the generator targets
@@ -424,6 +424,14 @@ func (g *treeGenerator) generateRewardsFile(treegen *rprewards.TreeGenerator) (*
424424 return treegen .GenerateTreeWithRuleset (g .ruleset )
425425}
426426
427+ func (g * treeGenerator ) serializeVotingPower (votingPowerFile * VotingPowerFile ) ([]byte , error ) {
428+ if g .prettyPrint {
429+ return json .MarshalIndent (votingPowerFile , "" , "\t " )
430+ }
431+
432+ return json .Marshal (votingPowerFile )
433+ }
434+
427435// Serializes the minipool performance file into JSON
428436func (g * treeGenerator ) serializeMinipoolPerformance (result * rprewards.GenerateTreeResult ) ([]byte , error ) {
429437 perfFile := result .MinipoolPerformanceFile
@@ -445,7 +453,7 @@ func (g *treeGenerator) serializeRewardsTree(rewardsFile rprewards.IRewardsFile)
445453}
446454
447455// Writes both the performance file and the rewards file to disk
448- func (g * treeGenerator ) writeFiles (result * rprewards.GenerateTreeResult ) error {
456+ func (g * treeGenerator ) writeFiles (result * rprewards.GenerateTreeResult , votingPowerFile * VotingPowerFile ) error {
449457 g .log .Printlnf ("Saving JSON files..." )
450458 rewardsFile := result .RewardsFile
451459 index := rewardsFile .GetIndex ()
@@ -489,6 +497,21 @@ func (g *treeGenerator) writeFiles(result *rprewards.GenerateTreeResult) error {
489497 }
490498
491499 g .log .Printlnf ("Saved rewards snapshot file to %s" , rewardsTreePath )
500+
501+ // Write the voting power file to disk
502+ if votingPowerFile != nil {
503+ votingPowerFileBytes , err := g .serializeVotingPower (votingPowerFile )
504+ if err != nil {
505+ return fmt .Errorf ("error serializing voting power file into JSON: %w" , err )
506+ }
507+
508+ votingFilePath := filepath .Join (g .outputDir , fmt .Sprintf ("rp-voting-power-%s-%d.json" , string (g .cfg .Smartnode .Network .Value .(cfgtypes.Network )), votingPowerFile .ConsensusSlot ))
509+ err = os .WriteFile (votingFilePath , votingPowerFileBytes , 0644 )
510+ if err != nil {
511+ return fmt .Errorf ("error saving voting power file to %s: %w" , votingFilePath , err )
512+ }
513+ g .log .Printlnf ("Saved voting power file to %s" , votingFilePath )
514+ }
492515 g .log .Printlnf ("Successfully generated rewards snapshot for interval %d" , index )
493516
494517 return nil
@@ -564,6 +587,7 @@ func (g *treeGenerator) approximateRethSpRewards() error {
564587
565588// Generate a complete rewards tree
566589func (g * treeGenerator ) generateTree () error {
590+ var votingPowerFile * VotingPowerFile
567591 args , err := g .getTreegenArgs ()
568592 if err != nil {
569593 return fmt .Errorf ("error compiling treegen arguments: %w" , err )
@@ -575,6 +599,12 @@ func (g *treeGenerator) generateTree() error {
575599 return err
576600 }
577601
602+ // If a voting power file was requested, generate it now.
603+ if g .generateVotingPower {
604+ votingPowerFile = g .GenerateVotingPower (args .state )
605+ votingPowerFile .Time = args .endTime
606+ }
607+
578608 // Generate the rewards file
579609 start := time .Now ()
580610 result , err := g .generateRewardsFile (treegen )
@@ -598,7 +628,7 @@ func (g *treeGenerator) generateTree() error {
598628 }
599629 }
600630
601- err = g .writeFiles (result )
631+ err = g .writeFiles (result , votingPowerFile )
602632 if err != nil {
603633 return err
604634 }
0 commit comments