stake_ci: add APIs for updating or removing individual contact info entries

tests

Author: ravyu-jump

src/disco/shred/fd_stake_ci.c

@@ -1,4 +1,5 @@
 #include "fd_stake_ci.h"
+#include "fd_shred_dest.h"
 #include "../../util/net/fd_ip4.h" /* Just for debug */
 
 #define SORT_NAME sort_pubkey
@@ -390,6 +391,100 @@ fd_stake_ci_set_identity( fd_stake_ci_t *     info,
   *info->identity_key = *identity_key;
 }
 
+void
+refresh_sdest( fd_stake_ci_t *            info,
+               fd_shred_dest_weighted_t * shred_dest_temp,
+               ulong                      cnt,
+               ulong                      staked_cnt,
+               fd_per_epoch_info_t *      ei ) {
+  sort_pubkey_inplace( shred_dest_temp + staked_cnt, cnt - staked_cnt );
+
+  fd_shred_dest_delete( fd_shred_dest_leave( ei->sdest ) );
+  ei->sdest = fd_shred_dest_join( fd_shred_dest_new( ei->_sdest, shred_dest_temp, cnt, ei->lsched, info->identity_key, ei->excluded_stake ) );
+  if( FD_UNLIKELY( ei->sdest==NULL ) ) {
+    FD_LOG_ERR(( "Too many validators have higher stake than this validator.  Cannot continue." ));
+  }
+}
+
+void
+ci_dest_add_one_unstaked( fd_stake_ci_t *            info,
+                          fd_shred_dest_weighted_t * new_entry,
+                          fd_per_epoch_info_t *      ei ) {
+  if( fd_shred_dest_cnt_all( ei->sdest )>=MAX_SHRED_DESTS ) {
+    FD_LOG_WARNING(( "Too many validators in shred table to add a new validator." ));
+  }
+  ulong cur_cnt = fd_shred_dest_cnt_all( ei->sdest );
+  for( ulong i=0UL; i<cur_cnt; i++ ) {
+    info->shred_dest_temp[ i ] = *fd_shred_dest_idx_to_dest( ei->sdest, (fd_shred_dest_idx_t)i );
+  }
+
+  /* TODO: Alternative batched copy using memcpy. Check with Philip if safe */
+  // fd_shred_dest_weighted_t * cur_dest = ei->sdest->all_destinations;
+  // fd_memcpy( info->shred_dest_temp, cur_dest, sizeof(fd_shred_dest_weighted_t)*cur_cnt );
+  info->shred_dest_temp[ cur_cnt++ ] = *new_entry;
+  refresh_sdest( info, info->shred_dest_temp, cur_cnt, fd_shred_dest_cnt_staked( ei->sdest ), ei );
+}
+
+void
+ci_dest_update_impl( fd_stake_ci_t *       info,
+                     fd_pubkey_t const *   pubkey,
+                     uint                  ip4,
+                     ushort                port,
+                     fd_per_epoch_info_t * ei ) {
+  fd_shred_dest_idx_t idx = fd_shred_dest_pubkey_to_idx( ei->sdest, pubkey );
+  if( idx==FD_SHRED_DEST_NO_DEST ) {
+    fd_shred_dest_weighted_t new_entry = { .pubkey = *pubkey, .ip4 = ip4, .port = port, .stake_lamports = 0UL };
+    ci_dest_add_one_unstaked( info, &new_entry, ei );
+    return;
+  }
+  fd_shred_dest_weighted_t * dest = fd_shred_dest_idx_to_dest( ei->sdest, idx );
+  dest->ip4                       = ip4;
+  dest->port                      = port;
+}
+
+void
+ci_dest_remove_impl( fd_stake_ci_t *       info,
+                     fd_pubkey_t const *   pubkey,
+                     fd_per_epoch_info_t * ei ) {
+  fd_shred_dest_idx_t idx = fd_shred_dest_pubkey_to_idx( ei->sdest, pubkey );
+  if( FD_UNLIKELY( idx==FD_SHRED_DEST_NO_DEST ) ) return;
+
+  fd_shred_dest_weighted_t * dest = fd_shred_dest_idx_to_dest( ei->sdest, idx );
+  if( FD_UNLIKELY( dest->stake_lamports>0UL ) ) {
+    /* A staked entry is not "removed", instead its "stale" address is
+       retained */
+    return;
+  }
+
+  ulong cur_cnt = fd_shred_dest_cnt_all( ei->sdest );
+  for( ulong i=0UL, j=0UL; i<cur_cnt; i++ ) {
+    if( FD_UNLIKELY( i==idx ) ) continue;
+    info->shred_dest_temp[ j++ ] = *fd_shred_dest_idx_to_dest( ei->sdest, (fd_shred_dest_idx_t) i );
+  }
+  /* TODO: Alternative batched copy using memcpy. Check with Philip if this is safe */
+  // fd_shred_dest_weighted_t * cur_dest = ei->sdest->all_destinations;
+  // fd_memcpy( info->shred_dest_temp, cur_dest, sizeof(fd_shred_dest_weighted_t)*(idx) );
+  // fd_memcpy( info->shred_dest_temp + idx, cur_dest + idx + 1UL, sizeof(fd_shred_dest_weighted_t)*(cur_cnt - idx - 1UL) );
+  refresh_sdest( info, info->shred_dest_temp, cur_cnt-1UL, fd_shred_dest_cnt_staked( ei->sdest ), ei );
+}
+
+void
+fd_stake_ci_dest_update( fd_stake_ci_t *       info,
+                         fd_pubkey_t const *   pubkey,
+                         uint                  ip4,
+                         ushort                port ) {
+  ci_dest_update_impl( info, pubkey, ip4, port, info->epoch_info+0UL );
+  ci_dest_update_impl( info, pubkey, ip4, port, info->epoch_info+1UL );
+}
+
+void
+fd_stake_ci_dest_remove( fd_stake_ci_t * info,
+                         fd_pubkey_t const * pubkey ) {
+  ci_dest_remove_impl( info, pubkey, info->epoch_info+0UL );
+  ci_dest_remove_impl( info, pubkey, info->epoch_info+1UL );
+
+}
+
 
 fd_shred_dest_t *
 fd_stake_ci_get_sdest_for_slot( fd_stake_ci_t const * info,

src/disco/shred/fd_stake_ci.h

@@ -93,7 +93,15 @@ fd_stake_ci_t * fd_stake_ci_join( void * mem );
 void * fd_stake_ci_leave ( fd_stake_ci_t * info );
 void * fd_stake_ci_delete( void          * mem  );
 
-/* fd_stake_ci_stake_msg_{init, fini} are used to handle messages
+/* Frankendancer and Firedancer's Gossip impls follow different regimes
+   for broadcasting Contact Infos. Firedancer employs an update-based
+   regime where we receive update/remove messages for individual contact
+   info entries. Frankendancer (and thusly Agave) performs a full table
+   broadcast. fd_stake_ci offers two sets of APIs that cater to the
+   different regimes. */
+
+/* Frankendancer only:
+   fd_stake_ci_stake_msg_{init, fini} are used to handle messages
    containing stake weight updates from the Rust side of the splice, and
    fd_stake_ci_dest_add_{init, fini} are used to handle messages
    containing contact info (potential shred destinations) updates from
@@ -152,6 +160,18 @@ void                       fd_stake_ci_stake_msg_fini( fd_stake_ci_t * info
 fd_shred_dest_weighted_t * fd_stake_ci_dest_add_init ( fd_stake_ci_t * info                                    );
 void                       fd_stake_ci_dest_add_fini ( fd_stake_ci_t * info, ulong                         cnt );
 
+/* Firedancer only:
+   The full client's Gossip update model publishes individual contact
+   info updates (update/insert or remove), which requires a different
+   set of dest_ APIs.
+
+   fd_stake_ci_dest_update updates (or adds, if necessary) a shred dest
+   entry. ip4 is in net order, port is in host order and are both
+   assumed to be non-zero. */
+
+void fd_stake_ci_dest_update( fd_stake_ci_t * info, fd_pubkey_t const * pubkey, uint ip4, ushort port );
+void fd_stake_ci_dest_remove( fd_stake_ci_t * info, fd_pubkey_t const * pubkey );
+
 
 /* fd_stake_ci_set_identity changes the identity of the locally running
    validator at runtime. */

src/disco/shred/test_stake_ci.c

@@ -445,6 +445,206 @@ test_staked_by_vote( void ) {
   fd_stake_ci_delete( fd_stake_ci_leave( info ) );
 }
 
+static void
+test_dest_update( void ) {
+  fd_stake_ci_t * info = fd_stake_ci_join( fd_stake_ci_new( _info, identity_key ) );
+
+  /* Set up initial state with only staked nodes */
+  fd_stake_ci_stake_msg_init( info, generate_stake_msg( stake_msg, 0UL, "ABC" ) );
+  fd_stake_ci_stake_msg_fini( info );
+  check_destinations( info, 0UL, "ABC", "I" );
+
+  /* Test updating existing staked node */
+  fd_pubkey_t pubkey_a;
+  memset( pubkey_a.uc, 'A', sizeof(fd_pubkey_t) );
+  fd_stake_ci_dest_update( info, &pubkey_a, 0x12345678U, 8080 );
+
+  fd_shred_dest_t * sdest = fd_stake_ci_get_sdest_for_slot( info, 0UL );
+  fd_shred_dest_idx_t idx_a = fd_shred_dest_pubkey_to_idx( sdest, &pubkey_a );
+  FD_TEST( idx_a != FD_SHRED_DEST_NO_DEST );
+  fd_shred_dest_weighted_t * dest_a = fd_shred_dest_idx_to_dest( sdest, idx_a );
+  FD_TEST( dest_a->ip4 == 0x12345678U );
+  FD_TEST( dest_a->port == 8080 );
+  FD_TEST( dest_a->stake_lamports > 0UL ); /* Should still be staked */
+
+  /* Test adding new unstaked node via update */
+  fd_pubkey_t pubkey_d;
+  memset( pubkey_d.uc, 'D', sizeof(fd_pubkey_t) );
+  fd_stake_ci_dest_update( info, &pubkey_d, 0x87654321U, 9090 );
+
+  /* D should now be in the unstaked list */
+  check_destinations( info, 0UL, "ABC", "DI" );
+
+  fd_shred_dest_idx_t idx_d = fd_shred_dest_pubkey_to_idx( sdest, &pubkey_d );
+  FD_TEST( idx_d != FD_SHRED_DEST_NO_DEST );
+  fd_shred_dest_weighted_t * dest_d = fd_shred_dest_idx_to_dest( sdest, idx_d );
+  FD_TEST( dest_d->ip4 == 0x87654321U );
+  FD_TEST( dest_d->port == 9090 );
+  FD_TEST( dest_d->stake_lamports == 0UL ); /* Should be unstaked */
+
+  /* Test adding multiple new unstaked nodes via update */
+  fd_pubkey_t pubkey_e, pubkey_f;
+  memset( pubkey_e.uc, 'E', sizeof(fd_pubkey_t) );
+  memset( pubkey_f.uc, 'F', sizeof(fd_pubkey_t) );
+
+  fd_stake_ci_dest_update( info, &pubkey_e, 0x11111111U, 1111 );
+  fd_stake_ci_dest_update( info, &pubkey_f, 0x22222222U, 2222 );
+
+  /* Check that E and F were added as unstaked */
+  check_destinations( info, 0UL, "ABC", "DEFI" );
+
+  /* Test updating an unstaked node's contact info */
+  fd_stake_ci_dest_update( info, &pubkey_d, 0x99999999U, 9999 );
+
+  idx_d = fd_shred_dest_pubkey_to_idx( sdest, &pubkey_d );
+  dest_d = fd_shred_dest_idx_to_dest( sdest, idx_d );
+  FD_TEST( dest_d->ip4 == 0x99999999U );
+  FD_TEST( dest_d->port == 9999 );
+  FD_TEST( dest_d->stake_lamports == 0UL ); /* Should still be unstaked */
+
+  /* Test that updates apply to both epochs */
+  fd_stake_ci_stake_msg_init( info, generate_stake_msg( stake_msg, 1UL, "AB" ) );
+  fd_stake_ci_stake_msg_fini( info );
+
+  /* Update should affect both epoch 0 and epoch 1 */
+  fd_pubkey_t pubkey_b;
+  memset( pubkey_b.uc, 'B', sizeof(fd_pubkey_t) );
+  fd_stake_ci_dest_update( info, &pubkey_b, 0x11223344U, 5555 );
+
+  /* Check epoch 0 */
+  fd_shred_dest_t * sdest0 = fd_stake_ci_get_sdest_for_slot( info, 0UL );
+  fd_shred_dest_idx_t idx_b0 = fd_shred_dest_pubkey_to_idx( sdest0, &pubkey_b );
+  fd_shred_dest_weighted_t * dest_b0 = fd_shred_dest_idx_to_dest( sdest0, idx_b0 );
+  FD_TEST( dest_b0->ip4 == 0x11223344U );
+  FD_TEST( dest_b0->port == 5555 );
+
+  /* Check epoch 1 */
+  fd_shred_dest_t * sdest1 = fd_stake_ci_get_sdest_for_slot( info, 1000UL );
+  fd_shred_dest_idx_t idx_b1 = fd_shred_dest_pubkey_to_idx( sdest1, &pubkey_b );
+  fd_shred_dest_weighted_t * dest_b1 = fd_shred_dest_idx_to_dest( sdest1, idx_b1 );
+  FD_TEST( dest_b1->ip4 == 0x11223344U );
+  FD_TEST( dest_b1->port == 5555 );
+
+  /* Test adding a new node that gets added to both epochs */
+  fd_pubkey_t pubkey_z;
+  memset( pubkey_z.uc, 'Z', sizeof(fd_pubkey_t) );
+  fd_stake_ci_dest_update( info, &pubkey_z, 0xAABBCCDDU, 7777 );
+
+  /* Z should be unstaked in both epochs */
+  check_destinations( info, 0UL, "ABC", "DEFIZ" );
+  /* C moved from staked to unstaked in epoch 1,
+     and with no contact info entry it should not
+     appear in the unstaked list. */
+  check_destinations( info, 1UL, "AB",  "DEFIZ" );
+
+  fd_pubkey_t pubkey_c;
+  memset( pubkey_c.uc, 'C', sizeof(fd_pubkey_t) );
+  fd_stake_ci_dest_update( info, &pubkey_c, 0xCCCCCCCCU, 8888 );
+
+  /* C should now be unstaked and present in epoch 1 */
+  check_destinations( info, 1UL, "AB", "CDEFIZ" );
+
+
+  fd_stake_ci_delete( fd_stake_ci_leave( info ) );
+}
+
+static void
+test_dest_remove( void ) {
+  fd_stake_ci_t * info = fd_stake_ci_join( fd_stake_ci_new( _info, identity_key ) );
+
+  /* Set up initial state with some staked and unstaked nodes */
+  fd_stake_ci_stake_msg_init( info, generate_stake_msg( stake_msg, 0UL, "ABC" ) );
+  fd_stake_ci_stake_msg_fini( info );
+
+  /* Build up destination list using only update operations */
+  fd_pubkey_t pubkey_a, pubkey_b, pubkey_c, pubkey_d, pubkey_e, pubkey_f;
+  memset( pubkey_a.uc, 'A', sizeof(fd_pubkey_t) );
+  memset( pubkey_b.uc, 'B', sizeof(fd_pubkey_t) );
+  memset( pubkey_c.uc, 'C', sizeof(fd_pubkey_t) );
+  memset( pubkey_d.uc, 'D', sizeof(fd_pubkey_t) );
+  memset( pubkey_e.uc, 'E', sizeof(fd_pubkey_t) );
+  memset( pubkey_f.uc, 'F', sizeof(fd_pubkey_t) );
+
+  /* Update staked nodes A, B, C with contact info */
+  fd_stake_ci_dest_update( info, &pubkey_a, 0x11111111U, 1111 );
+  fd_stake_ci_dest_update( info, &pubkey_b, 0x22222222U, 2222 );
+  fd_stake_ci_dest_update( info, &pubkey_c, 0x33333333U, 3333 );
+
+  /* Add unstaked nodes D, E, F via update */
+  fd_stake_ci_dest_update( info, &pubkey_d, 0x44444444U, 4444 );
+  fd_stake_ci_dest_update( info, &pubkey_e, 0x55555555U, 5555 );
+  fd_stake_ci_dest_update( info, &pubkey_f, 0x66666666U, 6666 );
+
+
+  check_destinations( info, 0UL, "ABC", "DEFI" );
+
+  /* Test removing unstaked node */
+  memset( pubkey_d.uc, 'D', sizeof(fd_pubkey_t) );
+  fd_stake_ci_dest_remove( info, &pubkey_d );
+
+  /* D should be removed from unstaked list */
+  check_destinations( info, 0UL, "ABC", "EFI" );
+
+  fd_shred_dest_t * sdest = fd_stake_ci_get_sdest_for_slot( info, 0UL );
+  fd_shred_dest_idx_t idx_d = fd_shred_dest_pubkey_to_idx( sdest, &pubkey_d );
+  FD_TEST( idx_d == FD_SHRED_DEST_NO_DEST ); /* Should not be found */
+
+  /* Test removing staked node (should NOT actually remove it, just clear contact info) */
+  memset( pubkey_a.uc, 'A', sizeof(fd_pubkey_t) );
+
+  /* First update A with some contact info */
+  fd_stake_ci_dest_update( info, &pubkey_a, 0x12345678U, 8080 );
+  fd_shred_dest_idx_t idx_a = fd_shred_dest_pubkey_to_idx( sdest, &pubkey_a );
+  fd_shred_dest_weighted_t * dest_a = fd_shred_dest_idx_to_dest( sdest, idx_a );
+  FD_TEST( dest_a->ip4 == 0x12345678U );
+  FD_TEST( dest_a->port == 8080 );
+
+  /* Now try to remove A - it should stay because it's staked */
+  fd_stake_ci_dest_remove( info, &pubkey_a );
+
+  /* A should still be in staked list */
+  check_destinations( info, 0UL, "ABC", "EFI" );
+  idx_a = fd_shred_dest_pubkey_to_idx( sdest, &pubkey_a );
+  FD_TEST( idx_a != FD_SHRED_DEST_NO_DEST ); /* Should still be found */
+  dest_a = fd_shred_dest_idx_to_dest( sdest, idx_a );
+  FD_TEST( dest_a->stake_lamports > 0UL ); /* Should still be staked */
+
+  /* Test removing non-existing node (should be no-op) */
+  fd_pubkey_t pubkey_z;
+  memset( pubkey_z.uc, 'Z', sizeof(fd_pubkey_t) );
+  fd_stake_ci_dest_remove( info, &pubkey_z );
+
+  /* Nothing should change */
+  check_destinations( info, 0UL, "ABC", "EFI" );
+
+  /* Test that removes apply to both epochs */
+  fd_stake_ci_stake_msg_init( info, generate_stake_msg( stake_msg, 1UL, "AB" ) );
+  fd_stake_ci_stake_msg_fini( info );
+
+  /* E should be unstaked in both epochs, remove it */
+  memset( pubkey_e.uc, 'E', sizeof(fd_pubkey_t) );
+  fd_stake_ci_dest_remove( info, &pubkey_e );
+
+  /* Check both epochs - E should be gone from both */
+  check_destinations( info, 0UL, "ABC", "FI" );
+  check_destinations( info, 1UL, "AB", "CFI" );
+
+  /* Verify E is not found in either epoch */
+  fd_shred_dest_t * sdest0 = fd_stake_ci_get_sdest_for_slot( info, 0UL );
+  fd_shred_dest_t * sdest1 = fd_stake_ci_get_sdest_for_slot( info, 1000UL );
+  FD_TEST( fd_shred_dest_pubkey_to_idx( sdest0, &pubkey_e ) == FD_SHRED_DEST_NO_DEST );
+  FD_TEST( fd_shred_dest_pubkey_to_idx( sdest1, &pubkey_e ) == FD_SHRED_DEST_NO_DEST );
+
+  /* Test removing multiple unstaked nodes */
+  memset( pubkey_f.uc, 'F', sizeof(fd_pubkey_t) );
+  fd_stake_ci_dest_remove( info, &pubkey_f );
+
+  check_destinations( info, 0UL, "ABC", "I" );
+  check_destinations( info, 1UL, "AB", "CI" );
+
+  fd_stake_ci_delete( fd_stake_ci_leave( info ) );
+}
+
 int
 main( int     argc,
       char ** argv ) {
@@ -473,6 +673,9 @@ main( int     argc,
   test_set_identity();
   test_staked_by_vote();
 
+  test_dest_update();
+  test_dest_remove();
+
   FD_LOG_NOTICE(( "pass" ));
   fd_halt();
   return 0;