Skip to content

Commit

Permalink
rebase and update to handkle tpu_vote_quic
Browse files Browse the repository at this point in the history
  • Loading branch information
gregcusack committed Nov 7, 2024
1 parent c635fee commit d88db4b
Show file tree
Hide file tree
Showing 2 changed files with 68 additions and 28 deletions.
32 changes: 15 additions & 17 deletions geyser-plugin-interface/src/geyser_plugin_interface.rs
Original file line number Diff line number Diff line change
Expand Up @@ -352,12 +352,10 @@ impl SlotStatus {
}
}

pub const PUBKEY_SIZE: usize = 32;

#[derive(Clone, Debug)]
#[repr(C)]
pub struct FfiPubkey {
pub pubkey: [u8; PUBKEY_SIZE],
pub pubkey: [u8; 32],
}

pub type Result<T> = std::result::Result<T, GeyserPluginError>;
Expand Down Expand Up @@ -463,6 +461,20 @@ pub trait GeyserPlugin: Any + Send + Sync + std::fmt::Debug {
Ok(())
}

/// Called when a ContactInfo is received from a node
#[allow(unused_variables)]
fn notify_node_update(&self, interface: &FfiContactInfoInterface) -> Result<()> {
Ok(())
}

/// Called when a node is removed from the network
/// TODO: may need to provide wrapper here? Also maybe ok
/// if we marshall to repr(c) for this...
#[allow(unused_variables)]
fn notify_node_removal(&self, pubkey: &FfiPubkey) -> Result<()> {
Ok(())
}

/// Check if the plugin is interested in account data
/// Default is true -- if the plugin is not interested in
/// account data, please return false.
Expand All @@ -484,20 +496,6 @@ pub trait GeyserPlugin: Any + Send + Sync + std::fmt::Debug {
false
}

/// Called when a message is received from a node
#[allow(unused_variables)]
fn notify_node_update(&self, interface: &FfiContactInfoInterface) -> Result<()> {
Ok(())
}

/// Called when a node is removed from the network
/// TODO: may need to provide wrapper here? Also maybe ok
/// if we marshall to repr(c) for this...
#[allow(unused_variables)]
fn notify_node_removal(&self, pubkey: &FfiPubkey) -> Result<()> {
Ok(())
}

/// Check if the plugin is interested in gossip m essages
/// Default is true -- if the plugin is not interested in
/// gossip messages, please return false.
Expand Down
64 changes: 53 additions & 11 deletions gossip/src/contact_info_ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -136,30 +136,38 @@ pub type ContactInfoGetShredVersionFn =
unsafe extern "C" fn(contact_info_ptr: ContactInfoPtr) -> u16;

/// Returns version of ContactInfo
/// TODO: figure out if we can fix this to not take in a *mut FfiVersion
/// although this may be the best bet. Returning a *const FfiVersion is
/// hard because we allocate it on the heap and we can't guarantee that
/// the caller will free it.
/// # Safety
/// - The ContactInfo pointer must be valid.
pub type ContactInfoGetVersionFn =
unsafe extern "C" fn(contact_info_ptr: ContactInfoPtr, ffi_version: *mut FfiVersion) -> bool;

/// Returns gossip address of ContactInfo
/// TODO: same as above
/// # Safety
/// - The ContactInfo pointer must be valid.
pub type ContactInfoGetGossipFn =
unsafe extern "C" fn(contact_info_ptr: ContactInfoPtr, socket: *mut FfiSocketAddr) -> bool;

/// Returns rpc address of ContactInfo
/// TODO: same as above
/// # Safety
/// - The ContactInfo pointer must be valid.
pub type ContactInfoGetRpcFn =
unsafe extern "C" fn(contact_info_ptr: ContactInfoPtr, socket: *mut FfiSocketAddr) -> bool;

/// Returns rpc_pubsub address of ContactInfo
/// TODO: same as above
/// # Safety
/// - The ContactInfo pointer must be valid.
pub type ContactInfoGetRpcPubsubpFn =
unsafe extern "C" fn(contact_info_ptr: ContactInfoPtr, socket: *mut FfiSocketAddr) -> bool;

/// Returns serve_repair address of ContactInfo
/// TODO: same as above
/// # Safety
/// - The ContactInfo pointer must be valid.
pub type ContactInfoGetServeRepairFn = unsafe extern "C" fn(
Expand All @@ -169,6 +177,7 @@ pub type ContactInfoGetServeRepairFn = unsafe extern "C" fn(
) -> bool;

/// Returns tpu address of ContactInfo
/// TODO: same as above
/// # Safety
/// - The ContactInfo pointer must be valid.
pub type ContactInfoGetTpuFn = unsafe extern "C" fn(
Expand All @@ -178,6 +187,7 @@ pub type ContactInfoGetTpuFn = unsafe extern "C" fn(
) -> bool;

/// Returns tpu_forwards address of ContactInfo
/// TODO: same as above
/// # Safety
/// - The ContactInfo pointer must be valid.
pub type ContactInfoGetTpuForwardsFn = unsafe extern "C" fn(
Expand All @@ -187,12 +197,17 @@ pub type ContactInfoGetTpuForwardsFn = unsafe extern "C" fn(
) -> bool;

/// Returns tpu_vote address of ContactInfo
/// TODO: same as above
/// # Safety
/// - The ContactInfo pointer must be valid.
pub type ContactInfoGetTpuVoteFn =
unsafe extern "C" fn(contact_info_ptr: ContactInfoPtr, socket: *mut FfiSocketAddr) -> bool;
pub type ContactInfoGetTpuVoteFn = unsafe extern "C" fn(
contact_info_ptr: ContactInfoPtr,
protocol: FfiProtocol,
socket: *mut FfiSocketAddr,
) -> bool;

/// Returns tvu address of ContactInfo
/// TODO: same as above
/// # Safety
/// - The ContactInfo pointer must be valid.
pub type ContactInfoGetTvuFn = unsafe extern "C" fn(
Expand Down Expand Up @@ -222,10 +237,6 @@ pub unsafe fn create_contact_info_interface(contact_info: &ContactInfo) -> FfiCo
contact_info.shred_version()
}

// TODO: figure out if we can fix this to not take in a *mut FfiVersion
// although this may be the best bet.
// Returning a *const FfiVersion is hard because we allocate it on the heap
// and we can't guarantee that the caller will free it.
extern "C" fn get_version(
contact_info_ptr: ContactInfoPtr,
ffi_version: *mut FfiVersion,
Expand Down Expand Up @@ -369,14 +380,17 @@ pub unsafe fn create_contact_info_interface(contact_info: &ContactInfo) -> FfiCo

extern "C" fn get_tpu_vote(
contact_info_ptr: ContactInfoPtr,
protocol: FfiProtocol,
socket: *mut FfiSocketAddr,
) -> bool {
if contact_info_ptr.is_null() || socket.is_null() {
return false;
}

let contact_info = unsafe { &*(contact_info_ptr as *const ContactInfo) };
match contact_info.tpu_vote() {
let protocol = Protocol::from(protocol); // Convert FfiProtocol to Protocol

match contact_info.tpu_vote(protocol) {
Ok(socket_addr) => {
let ffi_socket_addr = ffi_socket_addr_from_socket_addr(&socket_addr);
unsafe { *socket = ffi_socket_addr };
Expand Down Expand Up @@ -561,11 +575,15 @@ impl FfiContactInfoInterface {
}
}

pub fn tpu_vote(&self) -> Option<FfiSocketAddr> {
pub fn tpu_vote(&self, protocol: FfiProtocol) -> Option<FfiSocketAddr> {
let mut ffi_socket = FfiSocketAddr::default();

let success = unsafe {
(self.get_tpu_vote_fn)(self.contact_info_ptr, &mut ffi_socket as *mut FfiSocketAddr)
(self.get_tpu_vote_fn)(
self.contact_info_ptr,
protocol,
&mut ffi_socket as *mut FfiSocketAddr,
)
};

if success {
Expand Down Expand Up @@ -843,20 +861,44 @@ mod tests {

#[test]
fn test_get_tpu_vote() {
let contact_info = ContactInfo::new_localhost(&Pubkey::new_unique(), 123456789);
let mut contact_info = ContactInfo::new_localhost(&Pubkey::new_unique(), 123456789);
let interface = unsafe { create_contact_info_interface(&contact_info) };
// test udp
let mut ffi_socket = FfiSocketAddr::default();

let success = unsafe {
(interface.get_tpu_vote_fn)(
interface.contact_info_ptr,
FfiProtocol::UDP,
&mut ffi_socket as *mut FfiSocketAddr,
)
};

assert!(success);

let expected_socket_addr = contact_info.tpu_vote(Protocol::UDP).unwrap();
let actual_socket_addr = ffi_socket_addr_to_socket_addr(&ffi_socket);

assert_eq!(expected_socket_addr, actual_socket_addr);

// test quic
// TODO: remove once ContactInfo::new_localhost is updated to include set_tpu_vote_quic()
contact_info
.set_tpu_vote_quic((Ipv4Addr::LOCALHOST, 8009))
.unwrap();
let mut ffi_socket = FfiSocketAddr::default();

let success = unsafe {
(interface.get_tpu_vote_fn)(
interface.contact_info_ptr,
FfiProtocol::QUIC,
&mut ffi_socket as *mut FfiSocketAddr,
)
};

assert!(success);

let expected_socket_addr = contact_info.tpu_vote().unwrap();
let expected_socket_addr = contact_info.tpu_vote(Protocol::QUIC).unwrap();
let actual_socket_addr = ffi_socket_addr_to_socket_addr(&ffi_socket);

assert_eq!(expected_socket_addr, actual_socket_addr);
Expand Down

0 comments on commit d88db4b

Please sign in to comment.