rewards: cleaning up stake calculation

Author: ibhatt-jumptrading

src/flamenco/rewards/fd_rewards.c

@@ -1582,41 +1582,6 @@ fd_runtime_prepare_and_execute_txn( fd_banks_t *        banks,
 /* Epoch Boundary                                                             */
 /******************************************************************************/
 
-/* Update the epoch bank stakes cache with the delegated stake values from the slot bank cache.
-The slot bank cache will have been accumulating this epoch, and now we are at an epoch boundary
-we can safely update the epoch stakes cache with the latest values.
-
-In Solana, the stakes cache is updated after every transaction
-  (https://github.com/solana-labs/solana/blob/c091fd3da8014c0ef83b626318018f238f506435/runtime/src/bank.rs#L7587).
-As delegations have to warm up, the contents of the cache will not change inter-epoch. We can therefore update
-the cache only at epoch boundaries.
-
-https://github.com/solana-labs/solana/blob/c091fd3da8014c0ef83b626318018f238f506435/runtime/src/stakes.rs#L65 */
-static void
-fd_update_stake_delegations( fd_exec_slot_ctx_t * slot_ctx,
-                             fd_epoch_info_t *    temp_info ) {
-
-  fd_stake_delegations_t * stake_delegations = fd_stake_delegations_join( fd_bank_stake_delegations_locking_modify( slot_ctx->bank ) );
-
-  /* In one pass, iterate over all the new stake infos and insert the updated values into the epoch stakes cache
-      This assumes that there is enough memory pre-allocated for the stakes cache. */
-  for( ulong idx=temp_info->stake_infos_new_keys_start_idx; idx<temp_info->stake_infos_len; idx++ ) {
-    // Fetch and store the delegation associated with this stake account
-    fd_stake_delegations_update(
-        stake_delegations,
-        &temp_info->stake_infos[idx].account,
-        &temp_info->stake_infos[idx].stake.delegation.voter_pubkey,
-        temp_info->stake_infos[idx].stake.delegation.stake,
-        temp_info->stake_infos[idx].stake.delegation.activation_epoch,
-        temp_info->stake_infos[idx].stake.delegation.deactivation_epoch,
-        temp_info->stake_infos[idx].stake.credits_observed,
-        temp_info->stake_infos[idx].stake.delegation.warmup_cooldown_rate
-    );
-  }
-
-  fd_bank_stake_delegations_end_locking_modify( slot_ctx->bank );
-}
-
 /* Replace the stakes in T-2 (epoch_stakes) by the stakes at T-1 (next_epoch_stakes) */
 static void
 fd_update_epoch_stakes( fd_exec_slot_ctx_t * slot_ctx ) {
@@ -2249,12 +2214,11 @@ fd_runtime_process_new_epoch( fd_exec_slot_ctx_t * slot_ctx,
   }
 
   /* Updates stake history sysvar accumulated values. */
-  fd_stakes_activate_epoch( slot_ctx,
-                            new_rate_activation_epoch,
-                            &temp_info,
-                            runtime_spad );
-
-  fd_update_stake_delegations( slot_ctx, &temp_info );
+  fd_stakes_activate_epoch(
+      slot_ctx,
+      new_rate_activation_epoch,
+      &temp_info,
+      runtime_spad );
 
   /* Refresh vote accounts in stakes cache using updated stake weights, and merges slot bank vote accounts with the epoch bank vote accounts.
     https://github.com/anza-xyz/agave/blob/v2.1.6/runtime/src/stakes.rs#L363-L370 */

src/flamenco/runtime/fd_runtime.c

@@ -4,6 +4,7 @@
 #include "../runtime/context/fd_exec_slot_ctx.h"
 #include "../runtime/program/fd_stake_program.h"
 #include "../runtime/sysvar/fd_sysvar_stake_history.h"
+#include "fd_stake_delegations.h"
 
 ulong
 fd_stake_weights_by_node( fd_vote_accounts_global_t const * accs,
@@ -73,18 +74,14 @@ deserialize_and_update_vote_account( fd_exec_slot_ctx_t *
 }
 
 static void
-compute_stake_delegations(
-    fd_epoch_info_t *             temp_info,
-    ulong const                   epoch,
-    fd_stake_history_t const *    history,
-    ulong *                       new_rate_activation_epoch,
-    fd_stake_weight_t_mapnode_t * delegation_pool,
-    fd_stake_weight_t_mapnode_t * delegation_root,
-    ulong                         vote_states_pool_sz,
-    fd_spad_t *                   spad,
-    ulong                         end_idx
-) {
-  fd_epoch_info_pair_t const * stake_infos = temp_info->stake_infos;
+compute_stake_delegations( fd_bank_t *                   bank,
+                           ulong const                   epoch,
+                           fd_stake_history_t const *    history,
+                           ulong *                       new_rate_activation_epoch,
+                           fd_stake_weight_t_mapnode_t * delegation_pool,
+                           fd_stake_weight_t_mapnode_t * delegation_root,
+                           ulong                         vote_states_pool_sz,
+                           fd_spad_t *                   spad ) {
 
   FD_SPAD_FRAME_BEGIN( spad ) {
 
@@ -93,22 +90,43 @@ compute_stake_delegations(
   fd_stake_weight_t_mapnode_t * temp_pool = fd_stake_weight_t_map_join( fd_stake_weight_t_map_new( mem, vote_states_pool_sz ) );
   fd_stake_weight_t_mapnode_t * temp_root = NULL;
 
+  fd_stake_delegations_t const * stake_delegations = fd_bank_stake_delegations_locking_query( bank );
+  if( FD_UNLIKELY( !stake_delegations ) ) {
+    FD_LOG_CRIT(( "stake_delegations is NULL" ));
+  }
+  fd_stake_delegation_map_t * stake_delegation_map  = fd_stake_delegations_get_map( stake_delegations );
+  fd_stake_delegation_t *     stake_delegation_pool = fd_stake_delegations_get_pool( stake_delegations );
+
   fd_stake_weight_t_mapnode_t temp;
-  for( ulong i=0UL; i<end_idx; i++ ) {
-    fd_delegation_t const * delegation = &stake_infos[i].stake.delegation;
-    temp.elem.key = delegation->voter_pubkey;
+
+  for( fd_stake_delegation_map_iter_t iter = fd_stake_delegation_map_iter_init( stake_delegation_map, stake_delegation_pool );
+       !fd_stake_delegation_map_iter_done( iter, stake_delegation_map, stake_delegation_pool );
+       iter = fd_stake_delegation_map_iter_next( iter, stake_delegation_map, stake_delegation_pool ) ) {
+
+
+    fd_stake_delegation_t const * stake_delegation = fd_stake_delegation_map_iter_ele_const( iter, stake_delegation_map, stake_delegation_pool );
+
+    temp.elem.key = stake_delegation->vote_account;
 
     // Skip any delegations that are not in the delegation pool
     fd_stake_weight_t_mapnode_t * delegation_entry = fd_stake_weight_t_map_find( delegation_pool, delegation_root, &temp );
     if( FD_UNLIKELY( delegation_entry==NULL ) ) {
       continue;
     }
 
-    fd_stake_history_entry_t new_entry = fd_stake_activating_and_deactivating( delegation, epoch, history, new_rate_activation_epoch );
+    fd_delegation_t delegation = {
+      .voter_pubkey         = stake_delegation->vote_account,
+      .stake                = stake_delegation->stake,
+      .deactivation_epoch   = stake_delegation->deactivation_epoch,
+      .activation_epoch     = stake_delegation->activation_epoch,
+      .warmup_cooldown_rate = stake_delegation->warmup_cooldown_rate,
+    };
+
+    fd_stake_history_entry_t new_entry = fd_stake_activating_and_deactivating( &delegation, epoch, history, new_rate_activation_epoch );
     delegation_entry = fd_stake_weight_t_map_find( temp_pool, temp_root, &temp );
     if( FD_UNLIKELY( delegation_entry==NULL ) ) {
       delegation_entry = fd_stake_weight_t_map_acquire( temp_pool );
-      delegation_entry->elem.key   = delegation->voter_pubkey;
+      delegation_entry->elem.key   = stake_delegation->vote_account;
       delegation_entry->elem.stake = new_entry.effective;
       fd_stake_weight_t_map_insert( temp_pool, &temp_root, delegation_entry );
     } else {
@@ -124,6 +142,8 @@ compute_stake_delegations(
     output_delegation_node->elem.stake += elem->elem.stake;
   }
 
+  fd_bank_stake_delegations_end_locking_query( bank );
+
   } FD_SPAD_FRAME_END;
 
 }
@@ -189,16 +209,14 @@ fd_populate_vote_accounts( fd_exec_slot_ctx_t *       slot_ctx,
   fd_bank_vote_account_keys_end_locking_modify( slot_ctx->bank );
 
   compute_stake_delegations(
-      temp_info,
+      slot_ctx->bank,
       fd_bank_epoch_get( slot_ctx->bank ),
       history,
       new_rate_activation_epoch,
       pool,
       root,
       vote_states_pool_sz,
-      runtime_spad,
-      temp_info->stake_infos_len
-  );
+      runtime_spad );
 
   // Iterate over each vote account in the epoch stakes cache and populate the new vote accounts pool
   /* NOTE: we use epoch_bank->next_epoch_stakes because Agave indexes their epoch stakes cache by leader schedule epoch.
@@ -308,16 +326,14 @@ fd_refresh_vote_accounts( fd_exec_slot_ctx_t *       slot_ctx,
   }
 
   compute_stake_delegations(
-      temp_info,
+      slot_ctx->bank,
       fd_bank_epoch_get( slot_ctx->bank ),
       history,
       new_rate_activation_epoch,
       pool,
       root,
       vote_states_pool_sz,
-      runtime_spad,
-      temp_info->stake_infos_len
-  );
+      runtime_spad );
 
   // Iterate over each vote account in the epoch stakes cache and populate the new vote accounts pool
   ulong total_epoch_stake = 0UL;
@@ -403,13 +419,11 @@ fd_refresh_vote_accounts( fd_exec_slot_ctx_t *       slot_ctx,
 }
 
 static void
-accumulate_stake_cache_delegations(
-    fd_stake_delegations_t *   stake_delegations,
-    fd_stake_history_t const * history,
-    ulong *                    new_rate_activation_epoch,
-    fd_stake_history_entry_t * accumulator,
-    fd_epoch_info_t *          temp_info,
-    ulong                      epoch ) {
+accumulate_stake_cache_delegations( fd_stake_delegations_t *   stake_delegations,
+                                    fd_stake_history_t const * history,
+                                    ulong *                    new_rate_activation_epoch,
+                                    fd_stake_history_entry_t * accumulator,
+                                    ulong                      epoch ) {
 
   fd_stake_delegation_t *     pool = fd_stake_delegations_get_pool( stake_delegations );
   fd_stake_delegation_map_t * map  = fd_stake_delegations_get_map( stake_delegations );
@@ -423,19 +437,19 @@ accumulate_stake_cache_delegations(
        iter = fd_stake_delegation_map_iter_next( iter, map, pool ) ) {
     fd_stake_delegation_t const * stake_delegation = fd_stake_delegation_map_iter_ele_const( iter, map, pool );
 
-    ulong delegation_idx = temp_info->stake_infos_len++;
-
-    temp_info->stake_infos[delegation_idx].stake.credits_observed                = stake_delegation->credits_observed;
-    temp_info->stake_infos[delegation_idx].account                               = stake_delegation->stake_account;
-    temp_info->stake_infos[delegation_idx].stake.delegation.voter_pubkey         = stake_delegation->vote_account;
-    temp_info->stake_infos[delegation_idx].stake.delegation.stake                = stake_delegation->stake;
-    temp_info->stake_infos[delegation_idx].stake.delegation.activation_epoch     = stake_delegation->activation_epoch;
-    temp_info->stake_infos[delegation_idx].stake.delegation.deactivation_epoch   = stake_delegation->deactivation_epoch;
-    temp_info->stake_infos[delegation_idx].stake.delegation.warmup_cooldown_rate = stake_delegation->warmup_cooldown_rate;
-
-    fd_delegation_t * delegation = &temp_info->stake_infos[delegation_idx].stake.delegation;
-
-    fd_stake_history_entry_t new_entry = fd_stake_activating_and_deactivating( delegation, epoch, history, new_rate_activation_epoch );
+    fd_delegation_t delegation = {
+      .voter_pubkey         = stake_delegation->vote_account,
+      .stake                = stake_delegation->stake,
+      .activation_epoch     = stake_delegation->activation_epoch,
+      .deactivation_epoch   = stake_delegation->deactivation_epoch,
+      .warmup_cooldown_rate = stake_delegation->warmup_cooldown_rate,
+    };
+
+    fd_stake_history_entry_t new_entry = fd_stake_activating_and_deactivating(
+        &delegation,
+        epoch,
+        history,
+        new_rate_activation_epoch );
     effective    += new_entry.effective;
     activating   += new_entry.activating;
     deactivating += new_entry.deactivating;
@@ -456,7 +470,6 @@ fd_accumulate_stake_infos( fd_exec_slot_ctx_t const * slot_ctx,
                            fd_stake_history_t const * history,
                            ulong *                    new_rate_activation_epoch,
                            fd_stake_history_entry_t * accumulator,
-                           fd_epoch_info_t *          temp_info,
                            fd_spad_t *                runtime_spad ) {
 
   FD_SPAD_FRAME_BEGIN( runtime_spad ) {
@@ -468,9 +481,7 @@ fd_accumulate_stake_infos( fd_exec_slot_ctx_t const * slot_ctx,
       history,
       new_rate_activation_epoch,
       accumulator,
-      temp_info,
-      epoch
-  );
+      epoch );
 
   } FD_SPAD_FRAME_END;
 }
@@ -482,6 +493,7 @@ fd_stakes_activate_epoch( fd_exec_slot_ctx_t *  slot_ctx,
                           fd_epoch_info_t *     temp_info,
                           fd_spad_t *           runtime_spad ) {
 
+  (void)temp_info;
   fd_stake_delegations_t * stake_delegations = fd_stake_delegations_join( (void*)fd_bank_stake_delegations_locking_modify( slot_ctx->bank ) );
 
   /* Current stake delegations: list of all current delegations in stake_delegations
@@ -492,12 +504,6 @@ fd_stakes_activate_epoch( fd_exec_slot_ctx_t *  slot_ctx,
   fd_stake_history_t const * history = fd_sysvar_stake_history_read( slot_ctx->funk, slot_ctx->funk_txn, runtime_spad );
   if( FD_UNLIKELY( !history ) ) FD_LOG_ERR(( "StakeHistory sysvar is missing from sysvar cache" ));
 
-  ulong stake_delegations_size = fd_stake_delegations_cnt( stake_delegations );
-
-  temp_info->stake_infos_len = 0UL;
-  temp_info->stake_infos     = (fd_epoch_info_pair_t *)fd_spad_alloc( runtime_spad, FD_EPOCH_INFO_PAIR_ALIGN, sizeof(fd_epoch_info_pair_t)*stake_delegations_size );
-  fd_memset( temp_info->stake_infos, 0, sizeof(fd_epoch_info_pair_t)*stake_delegations_size );
-
   fd_stake_history_entry_t accumulator = {
     .effective    = 0UL,
     .activating   = 0UL,
@@ -511,7 +517,6 @@ fd_stakes_activate_epoch( fd_exec_slot_ctx_t *  slot_ctx,
       history,
       new_rate_activation_epoch,
       &accumulator,
-      temp_info,
       runtime_spad );
 
   /* https://github.com/anza-xyz/agave/blob/v2.1.6/runtime/src/stakes.rs#L359 */

src/flamenco/stakes/fd_stakes.c

@@ -70,7 +70,6 @@ fd_accumulate_stake_infos( fd_exec_slot_ctx_t const * slot_ctx,
                            fd_stake_history_t const * history,
                            ulong *                    new_rate_activation_epoch,
                            fd_stake_history_entry_t * accumulator,
-                           fd_epoch_info_t *          temp_info,
                            fd_spad_t *                runtime_spad );
 
 /* fd_store_stake_delegation is used to update fd_stake_delegations_t

src/flamenco/stakes/fd_stakes.h

@@ -3941,17 +3941,6 @@ void *fd_epoch_info_generate( void *mem, void **alloc_mem, fd_rng_t * rng ) {
   fd_epoch_info_t *self = (fd_epoch_info_t *) mem;
   *alloc_mem = (uchar *) *alloc_mem + sizeof(fd_epoch_info_t);
   fd_epoch_info_new(mem);
-  self->stake_infos_len = fd_rng_ulong( rng ) % 8;
-  if( self->stake_infos_len ) {
-    self->stake_infos = (fd_epoch_info_pair_t *) *alloc_mem;
-    *alloc_mem = (uchar *) *alloc_mem + sizeof(fd_epoch_info_pair_t)*self->stake_infos_len;
-    for( ulong i=0; i < self->stake_infos_len; i++ ) {
-      fd_epoch_info_pair_new( self->stake_infos + i );
-      fd_epoch_info_pair_generate( self->stake_infos + i, alloc_mem, rng );
-    }
-  } else {
-    self->stake_infos = NULL;
-  }
   ulong vote_states_len = fd_rng_ulong( rng ) % 8;
   self->vote_states_pool = fd_vote_info_pair_t_map_join_new( alloc_mem, vote_states_len );
   self->vote_states_root = NULL;

src/flamenco/types/fd_fuzz_types.h

@@ -25187,14 +25187,6 @@ ulong fd_vote_info_pair_size( fd_vote_info_pair_t const * self ) {
 
 int fd_epoch_info_encode( fd_epoch_info_t const * self, fd_bincode_encode_ctx_t * ctx ) {
   int err;
-  err = fd_bincode_uint64_encode( self->stake_infos_len, ctx );
-  if( FD_UNLIKELY(err) ) return err;
-  if( self->stake_infos_len ) {
-    for( ulong i=0; i < self->stake_infos_len; i++ ) {
-      err = fd_epoch_info_pair_encode( self->stake_infos + i, ctx );
-      if( FD_UNLIKELY( err ) ) return err;
-    }
-  }
   if( self->vote_states_root ) {
     ulong vote_states_len = fd_vote_info_pair_t_map_size( self->vote_states_pool, self->vote_states_root );
     err = fd_bincode_uint64_encode( vote_states_len, ctx );
@@ -25215,16 +25207,6 @@ int fd_epoch_info_encode( fd_epoch_info_t const * self, fd_bincode_encode_ctx_t
 static int fd_epoch_info_decode_footprint_inner( fd_bincode_decode_ctx_t * ctx, ulong * total_sz ) {
   if( ctx->data>=ctx->dataend ) { return FD_BINCODE_ERR_OVERFLOW; };
   int err = 0;
-  ulong stake_infos_len;
-  err = fd_bincode_uint64_decode( &stake_infos_len, ctx );
-  if( FD_UNLIKELY( err!=FD_BINCODE_SUCCESS ) ) return err;
-  if( stake_infos_len ) {
-    *total_sz += FD_EPOCH_INFO_PAIR_ALIGN + sizeof(fd_epoch_info_pair_t)*stake_infos_len;
-    for( ulong i=0; i < stake_infos_len; i++ ) {
-      err = fd_epoch_info_pair_decode_footprint_inner( ctx, total_sz );
-      if( FD_UNLIKELY( err!=FD_BINCODE_SUCCESS ) ) return err;
-    }
-  }
   ulong vote_states_len = 0UL;
   err = fd_bincode_uint64_decode( &vote_states_len, ctx );
   ulong vote_states_cnt = !!vote_states_len ? vote_states_len : 1;
@@ -25248,17 +25230,6 @@ int fd_epoch_info_decode_footprint( fd_bincode_decode_ctx_t * ctx, ulong * total
 }
 static void fd_epoch_info_decode_inner( void * struct_mem, void * * alloc_mem, fd_bincode_decode_ctx_t * ctx ) {
   fd_epoch_info_t * self = (fd_epoch_info_t *)struct_mem;
-  fd_bincode_uint64_decode_unsafe( &self->stake_infos_len, ctx );
-  if( self->stake_infos_len ) {
-    *alloc_mem = (void*)fd_ulong_align_up( (ulong)(*alloc_mem), FD_EPOCH_INFO_PAIR_ALIGN );
-    self->stake_infos = *alloc_mem;
-    *alloc_mem = (uchar *)(*alloc_mem) + sizeof(fd_epoch_info_pair_t)*self->stake_infos_len;
-    for( ulong i=0; i < self->stake_infos_len; i++ ) {
-      fd_epoch_info_pair_new( self->stake_infos + i );
-      fd_epoch_info_pair_decode_inner( self->stake_infos + i, alloc_mem, ctx );
-    }
-  } else
-    self->stake_infos = NULL;
   ulong vote_states_len;
   fd_bincode_uint64_decode_unsafe( &vote_states_len, ctx );
   self->vote_states_pool = fd_vote_info_pair_t_map_join_new( alloc_mem, vote_states_len );
@@ -25289,12 +25260,6 @@ void fd_epoch_info_new(fd_epoch_info_t * self) {
 void fd_epoch_info_walk( void * w, fd_epoch_info_t const * self, fd_types_walk_fn_t fun, const char *name, uint level, uint varint ) {
   (void) varint;
   fun( w, self, name, FD_FLAMENCO_TYPE_MAP, "fd_epoch_info", level++, 0 );
-  if( self->stake_infos_len ) {
-    fun( w, NULL, "stake_infos", FD_FLAMENCO_TYPE_ARR, "array", level++, 0 );
-    for( ulong i=0; i < self->stake_infos_len; i++ )
-      fd_epoch_info_pair_walk(w, self->stake_infos + i, fun, "epoch_info_pair", level, 0 );
-    fun( w, NULL, "stake_infos", FD_FLAMENCO_TYPE_ARR_END, "array", level--, 0 );
-  }
   if( self->vote_states_root ) {
     for( fd_vote_info_pair_t_mapnode_t * n = fd_vote_info_pair_t_map_minimum(self->vote_states_pool, self->vote_states_root ); n; n = fd_vote_info_pair_t_map_successor( self->vote_states_pool, n ) ) {
       fd_vote_info_pair_walk(w, &n->elem, fun, "vote_states", level, 0 );
@@ -25305,11 +25270,6 @@ void fd_epoch_info_walk( void * w, fd_epoch_info_t const * self, fd_types_walk_f
 }
 ulong fd_epoch_info_size( fd_epoch_info_t const * self ) {
   ulong size = 0;
-  do {
-    size += sizeof(ulong);
-    for( ulong i=0; i < self->stake_infos_len; i++ )
-      size += fd_epoch_info_pair_size( self->stake_infos + i );
-  } while(0);
   if( self->vote_states_root ) {
     size += sizeof(ulong);
     ulong max = fd_vote_info_pair_t_map_max( self->vote_states_pool );