refactor: move gRPC handlers to their own files

Author: rustybee42

mgmtd/src/grpc.rs

@@ -7,7 +7,7 @@ use crate::license::LicensedFeature;
 use crate::types::{ResolveEntityId, SqliteEnumExt};
 use anyhow::{Context as AContext, Result, anyhow, bail};
 use protobuf::{beegfs as pb, management as pm};
-use rusqlite::{OptionalExtension, Transaction, TransactionBehavior, params};
+use rusqlite::{OptionalExtension, Row, Transaction, TransactionBehavior, named_params, params};
 use shared::grpc::*;
 use shared::impl_grpc_handler;
 use shared::run_state::RunStateHandle;
@@ -21,13 +21,28 @@ use std::pin::Pin;
 use tonic::transport::{Identity, Server, ServerTlsConfig};
 use tonic::{Code, Request, Response, Status};
 
-mod buddy_group;
-mod license;
-mod misc;
-mod node;
-mod pool;
-mod quota;
-mod target;
+mod common;
+
+mod assign_pool;
+mod create_buddy_group;
+mod create_pool;
+mod delete_buddy_group;
+mod delete_node;
+mod delete_pool;
+mod delete_target;
+mod get_buddy_groups;
+mod get_license;
+mod get_nodes;
+mod get_pools;
+mod get_quota_limits;
+mod get_quota_usage;
+mod get_targets;
+mod mirror_root_inode;
+mod set_alias;
+mod set_default_quota_limits;
+mod set_quota_limits;
+mod set_target_state;
+mod start_resync;
 
 /// Management gRPC service implementation struct
 #[derive(Debug)]
@@ -44,8 +59,9 @@ impl pm::management_server::Management for ManagementService {
     // Example: Implement pm::management_server::Management::set_alias using the impl_grpc_handler
     // macro
     impl_grpc_handler! {
-        // <the function to implement (as defined by the trait)> => <the actual, custom handler function to call>,
-        set_alias => misc::set_alias,
+        // the function to implement (as defined by the trait) as well as the handler to call (must
+        // be named the same and in a submodule named the same),
+        set_alias,
         // <request message passed to the fn impl (as defined by the trait)> => <response message,
         // returned by the fn impl (as defined by the trait)>,
         pm::SetAliasRequest => pm::SetAliasResponse,
@@ -54,102 +70,102 @@ impl pm::management_server::Management for ManagementService {
     }
 
     impl_grpc_handler! {
-        get_nodes => node::get,
+        get_nodes,
         pm::GetNodesRequest => pm::GetNodesResponse,
         "Get nodes"
     }
     impl_grpc_handler! {
-        delete_node => node::delete,
+        delete_node,
         pm::DeleteNodeRequest => pm::DeleteNodeResponse,
         "Delete node"
     }
 
     impl_grpc_handler! {
-        get_targets => target::get,
+        get_targets,
         pm::GetTargetsRequest => pm::GetTargetsResponse,
         "Get targets"
     }
     impl_grpc_handler! {
-        delete_target => target::delete,
+        delete_target,
         pm::DeleteTargetRequest => pm::DeleteTargetResponse,
         "Delete target"
     }
     impl_grpc_handler! {
-        set_target_state => target::set_state,
+        set_target_state,
         pm::SetTargetStateRequest => pm::SetTargetStateResponse,
         "Set target state"
     }
 
     impl_grpc_handler! {
-        get_pools => pool::get,
+        get_pools,
         pm::GetPoolsRequest => pm::GetPoolsResponse,
         "Get pools"
     }
     impl_grpc_handler! {
-        create_pool => pool::create,
+        create_pool,
         pm::CreatePoolRequest => pm::CreatePoolResponse,
         "Create pool"
     }
     impl_grpc_handler! {
-        assign_pool => pool::assign,
+        assign_pool,
         pm::AssignPoolRequest => pm::AssignPoolResponse,
         "Assign pool"
     }
     impl_grpc_handler! {
-        delete_pool => pool::delete,
+        delete_pool,
         pm::DeletePoolRequest => pm::DeletePoolResponse,
         "Delete pool"
     }
 
     impl_grpc_handler! {
-        get_buddy_groups => buddy_group::get,
+        get_buddy_groups,
         pm::GetBuddyGroupsRequest => pm::GetBuddyGroupsResponse,
         "Get buddy groups"
     }
     impl_grpc_handler! {
-        create_buddy_group => buddy_group::create,
+        create_buddy_group,
         pm::CreateBuddyGroupRequest => pm::CreateBuddyGroupResponse,
         "Create buddy group"
     }
     impl_grpc_handler! {
-        delete_buddy_group => buddy_group::delete,
+        delete_buddy_group,
         pm::DeleteBuddyGroupRequest => pm::DeleteBuddyGroupResponse,
         "Delete buddy group"
     }
     impl_grpc_handler! {
-        mirror_root_inode => buddy_group::mirror_root_inode,
+        mirror_root_inode,
         pm::MirrorRootInodeRequest => pm::MirrorRootInodeResponse,
         "Mirror root inode"
     }
     impl_grpc_handler! {
-        start_resync => buddy_group::start_resync,
+        start_resync,
         pm::StartResyncRequest => pm::StartResyncResponse,
         "Start resync"
     }
 
     impl_grpc_handler! {
-        set_default_quota_limits => quota::set_default_quota_limits,
+        set_default_quota_limits,
         pm::SetDefaultQuotaLimitsRequest => pm::SetDefaultQuotaLimitsResponse,
         "Set default quota limits"
     }
     impl_grpc_handler! {
-        set_quota_limits => quota::set_quota_limits,
+        set_quota_limits,
         pm::SetQuotaLimitsRequest => pm::SetQuotaLimitsResponse,
         "Set quota limits"
     }
     impl_grpc_handler! {
-        get_quota_limits => quota::get_quota_limits,
+        get_quota_limits,
         pm::GetQuotaLimitsRequest => STREAM(GetQuotaLimitsStream, pm::GetQuotaLimitsResponse),
         "Get quota limits"
     }
     impl_grpc_handler! {
-        get_quota_usage => quota::get_quota_usage,
+        get_quota_usage,
         pm::GetQuotaUsageRequest => STREAM(GetQuotaUsageStream, pm::GetQuotaUsageResponse),
         "Get quota usage"
     }
 
     impl_grpc_handler! {
-        get_license => license::get,
+        get_license,
         pm::GetLicenseRequest => pm::GetLicenseResponse,
         "Get license"
     }

mgmtd/src/grpc/assign_pool.rs

@@ -0,0 +1,95 @@
+use super::*;
+use shared::bee_msg::storage_pool::RefreshStoragePools;
+
+/// Assigns a pool to a list of targets and buddy groups.
+pub(crate) async fn assign_pool(
+    ctx: Context,
+    req: pm::AssignPoolRequest,
+) -> Result<pm::AssignPoolResponse> {
+    needs_license(&ctx, LicensedFeature::Storagepool)?;
+    fail_on_pre_shutdown(&ctx)?;
+
+    let pool: EntityId = required_field(req.pool)?.try_into()?;
+
+    let pool = ctx
+        .db
+        .write_tx(move |tx| {
+            let pool = pool.resolve(tx, EntityType::Pool)?;
+            do_assign(tx, pool.num_id().try_into()?, req.targets, req.buddy_groups)?;
+            Ok(pool)
+        })
+        .await?;
+
+    log::info!("Pool assigned: {pool}");
+
+    notify_nodes(
+        &ctx,
+        &[NodeType::Meta, NodeType::Storage],
+        &RefreshStoragePools { ack_id: "".into() },
+    )
+    .await;
+
+    Ok(pm::AssignPoolResponse {
+        pool: Some(pool.into()),
+    })
+}
+
+/// Do the actual assign work
+pub(super) fn do_assign(
+    tx: &Transaction,
+    pool_id: PoolId,
+    targets: Vec<pb::EntityIdSet>,
+    groups: Vec<pb::EntityIdSet>,
+) -> Result<()> {
+    // Target being part of a buddy group can not be assigned individually
+    let mut check_group_membership = tx.prepare_cached(sql!(
+        "SELECT COUNT(*) FROM storage_buddy_groups AS g
+        INNER JOIN targets AS p_t ON p_t.target_id = g.p_target_id AND p_t.node_type = g.node_type
+        INNER JOIN targets AS s_t ON s_t.target_id = g.s_target_id AND s_t.node_type = g.node_type
+        WHERE p_t.target_uid = ?1 OR s_t.target_uid = ?1"
+    ))?;
+
+    let mut assign_target = tx.prepare_cached(sql!(
+        "UPDATE targets SET pool_id = ?1 WHERE target_uid = ?2"
+    ))?;
+
+    // Do the checks and assign for each target in the given list. This is expensive, but shouldn't
+    // matter as this command should only be run occasionally and not with very large lists of
+    // targets.
+    for t in targets {
+        let eid = EntityId::try_from(t)?;
+        let target = eid.resolve(tx, EntityType::Target)?;
+        if check_group_membership.query_row([target.uid], |row| row.get::<_, i64>(0))? > 0 {
+            bail!("Target {eid} can't be assigned directly as it's part of a buddy group");
+        }
+
+        assign_target.execute(params![pool_id, target.uid])?;
+    }
+
+    let mut assign_group = tx.prepare_cached(sql!(
+        "UPDATE buddy_groups SET pool_id = ?1 WHERE group_uid = ?2"
+    ))?;
+
+    // Targets being part of buddy groups are auto-assigned to the new pool
+    let mut assign_grouped_targets = tx.prepare_cached(sql!(
+        "UPDATE targets SET pool_id = ?1
+        FROM (
+            SELECT p_t.target_uid AS p_uid, s_t.target_uid AS s_uid FROM buddy_groups AS g
+            INNER JOIN targets AS p_t ON p_t.target_id = g.p_target_id AND p_t.node_type = g.node_type
+            INNER JOIN targets AS s_t ON s_t.target_id = g.s_target_id AND s_t.node_type = g.node_type
+            WHERE group_uid = ?2
+        )
+        WHERE target_uid IN (p_uid, s_uid)"
+    ))?;
+
+    // Assign each group and their targets to the new pool
+    for g in groups {
+        let eid = EntityId::try_from(g)?;
+        let group = eid.resolve(tx, EntityType::BuddyGroup)?;
+
+        assign_group.execute(params![pool_id, group.uid])?;
+        assign_grouped_targets.execute(params![pool_id, group.uid])?;
+    }
+
+    Ok(())
+}