feat: add reload_plugin and unload_plugin admin RPCs

Author: 0xsouravm

crates/core/src/lib.rs

@@ -54,6 +54,8 @@ pub async fn start_local_surfnet(
 #[derive(Debug)]
 pub enum PluginManagerCommand {
     LoadConfig(Uuid, PluginConfig, Sender<String>),
+    UnloadPlugin(Uuid, Sender<Result<(), String>>),
+    ReloadPlugin(Uuid, PluginConfig, Sender<Result<String, String>>),
 }
 
 #[cfg(test)]

crates/core/src/rpc/admin.rs

@@ -770,15 +770,71 @@ impl AdminRpc for SurfpoolAdminRpc {
 
     fn reload_plugin(
         &self,
-        _meta: Self::Metadata,
-        _name: String,
-        _config_file: String,
+        meta: Self::Metadata,
+        name: String,
+        config_file: String,
     ) -> BoxFuture<Result<()>> {
-        not_implemented_err_async("reload_plugin")
+        // Parse the UUID from the name parameter
+        let uuid = match Uuid::parse_str(&name) {
+            Ok(uuid) => uuid,
+            Err(e) => return Box::pin(async move {
+                Err(jsonrpc_core::Error::invalid_params(format!("Invalid UUID: {}", e)))
+            }),
+        };
+
+        // Parse the new configuration
+        let config = match serde_json::from_str::<PluginConfig>(&config_file)
+            .map_err(|e| format!("failed to deserialize plugin config: {e}"))
+        {
+            Ok(config) => config,
+            Err(e) => return Box::pin(async move { Err(jsonrpc_core::Error::invalid_params(&e)) }),
+        };
+
+        let Some(ctx) = meta else {
+            return Box::pin(async move { Err(jsonrpc_core::Error::internal_error()) });
+        };
+
+        let (tx, rx) = crossbeam_channel::bounded(1);
+        let _ = ctx
+            .plugin_manager_commands_tx
+            .send(PluginManagerCommand::ReloadPlugin(uuid, config, tx));
+
+        let Ok(result) = rx.recv_timeout(Duration::from_secs(10)) else {
+            return Box::pin(async move { Err(jsonrpc_core::Error::internal_error()) });
+        };
+
+        Box::pin(async move {
+            result
+                .map(|_| ())  // Convert Ok(String) to Ok(())
+                .map_err(|e| jsonrpc_core::Error::invalid_params(&e))
+        })
     }
 
-    fn unload_plugin(&self, _meta: Self::Metadata, _name: String) -> BoxFuture<Result<()>> {
-        not_implemented_err_async("unload_plugin")
+    fn unload_plugin(&self, meta: Self::Metadata, name: String) -> BoxFuture<Result<()>> {
+        // Parse the UUID from the name parameter
+        let uuid = match Uuid::parse_str(&name) {
+            Ok(uuid) => uuid,
+            Err(e) => return Box::pin(async move {
+                Err(jsonrpc_core::Error::invalid_params(format!("Invalid UUID: {}", e)))
+            }),
+        };
+
+        let Some(ctx) = meta else {
+            return Box::pin(async move { Err(jsonrpc_core::Error::internal_error()) });
+        };
+
+        let (tx, rx) = crossbeam_channel::bounded(1);
+        let _ = ctx
+            .plugin_manager_commands_tx
+            .send(PluginManagerCommand::UnloadPlugin(uuid, tx));
+
+        let Ok(result) = rx.recv_timeout(Duration::from_secs(10)) else {
+            return Box::pin(async move { Err(jsonrpc_core::Error::internal_error()) });
+        };
+
+        Box::pin(async move {
+            result.map_err(|e| jsonrpc_core::Error::invalid_params(&e))
+        })
     }
 
     fn load_plugin(&self, meta: Self::Metadata, config_file: String) -> BoxFuture<Result<String>> {
@@ -798,7 +854,9 @@ impl AdminRpc for SurfpoolAdminRpc {
         let Ok(endpoint_url) = rx.recv_timeout(Duration::from_secs(10)) else {
             return Box::pin(async move { Err(jsonrpc_core::Error::internal_error()) });
         };
-        Box::pin(async move { Ok(endpoint_url) })
+        // Return a JSON string containing both UUID and endpoint URL
+        let response = format!(r#"{{"uuid": "{}", "endpoint": "{}"}}"#, uuid, endpoint_url);
+        Box::pin(async move { Ok(response) })
     }
 
     fn list_plugins(&self, _meta: Self::Metadata) -> BoxFuture<Result<Vec<String>>> {

crates/core/src/runloops/mod.rs

@@ -314,6 +314,10 @@ fn start_geyser_runloop(
 
         let mut surfpool_plugin_manager: Vec<Box<dyn GeyserPlugin>> = vec![];
 
+        // Map between each plugin's UUID to its position (index) in the surfpool_plugin_manager Vec.
+        // Allows for easier reload/unload
+        let mut plugin_uuid_map: HashMap<crate::Uuid, usize> = HashMap::new();
+
         #[cfg(feature = "geyser_plugin")]
         for plugin_config_path in plugin_config_paths.into_iter() {
             let plugin_manifest_location = FileLocation::from_path(plugin_config_path);
@@ -367,6 +371,92 @@ fn start_geyser_runloop(
 
         let ipc_router = RouterProxy::new();
 
+        // Helper function to load a subgraph plugin
+        #[cfg(feature = "subgraph")]
+        let load_subgraph_plugin = |uuid: uuid::Uuid,
+                                      config: txtx_addon_network_svm_types::subgraph::PluginConfig,
+                                      surfpool_plugin_manager: &mut Vec<Box<dyn GeyserPlugin>>,
+                                      plugin_uuid_map: &mut HashMap<uuid::Uuid, usize>,
+                                      indexing_enabled: &mut bool|
+         -> Result<String, String> {
+            let _ = subgraph_commands_tx.send(SubgraphCommand::CreateCollection(
+                uuid,
+                config.data.clone(),
+                crossbeam_channel::bounded(0).0,  // Temporary sender, will be replaced
+            ));
+            let mut plugin = SurfpoolSubgraphPlugin::default();
+
+            let (server, ipc_token) =
+                IpcOneShotServer::<IpcReceiver<DataIndexingCommand>>::new()
+                    .expect("Failed to create IPC one-shot server.");
+            let subgraph_plugin_config = SubgraphPluginConfig {
+                uuid,
+                ipc_token,
+                subgraph_request: config.data.clone(),
+            };
+
+            let config_file = serde_json::to_string(&subgraph_plugin_config)
+                .map_err(|e| format!("Failed to serialize subgraph plugin config: {:?}", e))?;
+
+            plugin
+                .on_load(&config_file, false)
+                .map_err(|e| format!("Failed to load Geyser plugin: {:?}", e))?;
+
+            if let Ok((_, rx)) = server.accept() {
+                let subgraph_rx = ipc_router
+                    .route_ipc_receiver_to_new_crossbeam_receiver::<DataIndexingCommand>(rx);
+                let _ = subgraph_commands_tx.send(SubgraphCommand::ObserveCollection(subgraph_rx));
+            };
+
+            *indexing_enabled = true;
+
+            let plugin: Box<dyn GeyserPlugin> = Box::new(plugin);
+            let plugin_index = surfpool_plugin_manager.len();
+            surfpool_plugin_manager.push(plugin);
+            plugin_uuid_map.insert(uuid, plugin_index);
+
+            let _ = simnet_events_tx.send(SimnetEvent::PluginLoaded("surfpool-subgraph".into()));
+            Ok(format!("http://localhost:8899/subgraph/{}", uuid))  // Return endpoint URL
+        };
+
+        // Helper function to unload a plugin by UUID
+        #[cfg(feature = "subgraph")]
+        let unload_plugin_by_uuid = |uuid: uuid::Uuid,
+                                       surfpool_plugin_manager: &mut Vec<Box<dyn GeyserPlugin>>,
+                                       plugin_uuid_map: &mut HashMap<uuid::Uuid, usize>,
+                                       indexing_enabled: &mut bool|
+         -> Result<(), String> {
+            let plugin_index = *plugin_uuid_map
+                .get(&uuid)
+                .ok_or_else(|| format!("Plugin {} not found", uuid))?;
+
+            if plugin_index >= surfpool_plugin_manager.len() {
+                return Err(format!("Plugin index {} out of bounds", plugin_index));
+            }
+
+            // Call on_unload before removing
+            surfpool_plugin_manager[plugin_index].on_unload();
+
+            // Remove the plugin from the list
+            surfpool_plugin_manager.remove(plugin_index);
+            plugin_uuid_map.remove(&uuid);
+
+            // Update all UUIDs that had indices after the removed one
+            for (_, idx) in plugin_uuid_map.iter_mut() {
+                if *idx > plugin_index {
+                    *idx -= 1;
+                }
+            }
+
+            // Check if we should disable indexing
+            if surfpool_plugin_manager.is_empty() {
+                *indexing_enabled = false;
+            }
+
+            let _ = simnet_events_tx.send(SimnetEvent::info(format!("Plugin {} unloaded", uuid)));
+            Ok(())
+        };
+
         let err = loop {
             use agave_geyser_plugin_interface::geyser_plugin_interface::{ReplicaAccountInfoV3, ReplicaAccountInfoVersions};
 
@@ -383,37 +473,46 @@ fn start_geyser_runloop(
                                 }
                                 #[cfg(feature = "subgraph")]
                                 PluginManagerCommand::LoadConfig(uuid, config, notifier) => {
-                                    let _ = subgraph_commands_tx.send(SubgraphCommand::CreateCollection(uuid, config.data.clone(), notifier));
-                                    let mut plugin = SurfpoolSubgraphPlugin::default();
-
-                                    let (server, ipc_token) = IpcOneShotServer::<IpcReceiver<DataIndexingCommand>>::new().expect("Failed to create IPC one-shot server.");
-                                    let subgraph_plugin_config = SubgraphPluginConfig {
-                                        uuid,
-                                        ipc_token,
-                                        subgraph_request: config.data.clone()
-                                    };
-
-                                    let config_file = match serde_json::to_string(&subgraph_plugin_config) {
-                                        Ok(c) => c,
+                                    match load_subgraph_plugin(uuid, config, &mut surfpool_plugin_manager, &mut plugin_uuid_map, &mut indexing_enabled) {
+                                        Ok(endpoint_url) => {
+                                            let _ = notifier.send(endpoint_url);
+                                        }
                                         Err(e) => {
-                                            let _ = simnet_events_tx.send(SimnetEvent::error(format!("Failed to serialize subgraph plugin config: {:?}", e)));
-                                            continue;
+                                            let _ = simnet_events_tx.send(SimnetEvent::error(format!("Failed to load plugin: {}", e)));
                                         }
-                                    };
-
-                                    if let Err(e) = plugin.on_load(&config_file, false) {
-                                        let _ = simnet_events_tx.send(SimnetEvent::error(format!("Failed to load Geyser plugin: {:?}", e)));
-                                    };
-                                    if let Ok((_, rx)) = server.accept() {
-                                        let subgraph_rx = ipc_router.route_ipc_receiver_to_new_crossbeam_receiver::<DataIndexingCommand>(rx);
-                                        let _ = subgraph_commands_tx.send(SubgraphCommand::ObserveCollection(subgraph_rx));
-                                    };
-
-                                    indexing_enabled = true;
+                                    }
+                                }
+                                #[cfg(not(feature = "subgraph"))]
+                                PluginManagerCommand::UnloadPlugin(_, _) => {
+                                    continue;
+                                }
+                                #[cfg(feature = "subgraph")]
+                                PluginManagerCommand::UnloadPlugin(uuid, notifier) => {
+                                    let result = unload_plugin_by_uuid(uuid, &mut surfpool_plugin_manager, &mut plugin_uuid_map, &mut indexing_enabled);
+                                    let _ = notifier.send(result);
+                                }
+                                #[cfg(not(feature = "subgraph"))]
+                                PluginManagerCommand::ReloadPlugin(_, _, _) => {
+                                    continue;
+                                }
+                                #[cfg(feature = "subgraph")]
+                                PluginManagerCommand::ReloadPlugin(uuid, config, notifier) => {
+                                    // First, unload the old plugin
+                                    if let Err(e) = unload_plugin_by_uuid(uuid, &mut surfpool_plugin_manager, &mut plugin_uuid_map, &mut indexing_enabled) {
+                                        let _ = notifier.send(Err(e));
+                                        continue;
+                                    }
 
-                                    let plugin: Box<dyn GeyserPlugin> = Box::new(plugin);
-                                    surfpool_plugin_manager.push(plugin);
-                                    let _ = simnet_events_tx.send(SimnetEvent::PluginLoaded("surfpool-subgraph".into()));
+                                    // Then, load the new plugin with the same UUID
+                                    match load_subgraph_plugin(uuid, config, &mut surfpool_plugin_manager, &mut plugin_uuid_map, &mut indexing_enabled) {
+                                        Ok(endpoint_url) => {
+                                            let _ = notifier.send(Ok(endpoint_url));
+                                            let _ = simnet_events_tx.send(SimnetEvent::info(format!("Plugin {} reloaded", uuid)));
+                                        }
+                                        Err(e) => {
+                                            let _ = notifier.send(Err(e));
+                                        }
+                                    }
                                 }
                             }
                         },