From afe30dc97df97283efafe46ed50d4486b6be926f Mon Sep 17 00:00:00 2001 From: Rusty Bee <145002912+rustybee42@users.noreply.github.com> Date: Wed, 12 Aug 2026 16:10:44 +0200 Subject: [PATCH 1/9] refactor: make unwrapping required and optional protobuf message fields easier --- mgmtd/src/grpc/assign_pool.rs | 2 +- mgmtd/src/grpc/create_buddy_group.rs | 8 ++--- mgmtd/src/grpc/create_pool.rs | 4 +-- mgmtd/src/grpc/delete_buddy_group.rs | 2 +- mgmtd/src/grpc/delete_node.rs | 2 +- mgmtd/src/grpc/delete_pool.rs | 2 +- mgmtd/src/grpc/delete_target.rs | 2 +- mgmtd/src/grpc/set_alias.rs | 2 +- mgmtd/src/grpc/set_default_quota_limits.rs | 2 +- mgmtd/src/grpc/set_quota_limits.rs | 4 +-- mgmtd/src/grpc/set_target_state.rs | 2 +- mgmtd/src/grpc/start_resync.rs | 2 +- shared/src/grpc.rs | 37 +++++++++++++++++++--- shared/src/impl_macros.rs | 13 +++++++- 14 files changed, 61 insertions(+), 23 deletions(-) diff --git a/mgmtd/src/grpc/assign_pool.rs b/mgmtd/src/grpc/assign_pool.rs index 6b9e968b..82238a5f 100644 --- a/mgmtd/src/grpc/assign_pool.rs +++ b/mgmtd/src/grpc/assign_pool.rs @@ -9,7 +9,7 @@ pub(crate) async fn assign_pool( fail_on_missing_license(app, LicensedFeature::Storagepool)?; fail_on_pre_shutdown(app)?; - let pool: EntityId = required_field(req.pool)?.try_into()?; + let pool: EntityId = required_field(req.pool)?; let pool = app .write_tx(move |tx| { diff --git a/mgmtd/src/grpc/create_buddy_group.rs b/mgmtd/src/grpc/create_buddy_group.rs index cd561193..a7cf97e4 100644 --- a/mgmtd/src/grpc/create_buddy_group.rs +++ b/mgmtd/src/grpc/create_buddy_group.rs @@ -11,10 +11,10 @@ pub(crate) async fn create_buddy_group( fail_on_pre_shutdown(app)?; let node_type: NodeTypeServer = req.node_type().try_into()?; - let alias: Alias = required_field(req.alias)?.try_into()?; - let num_id: BuddyGroupId = req.num_id.unwrap_or_default().try_into()?; - let p_target: EntityId = required_field(req.primary_target)?.try_into()?; - let s_target: EntityId = required_field(req.secondary_target)?.try_into()?; + let alias: Alias = required_field(req.alias)?; + let num_id: BuddyGroupId = optional_field(req.num_id)?.unwrap_or_default(); + let p_target: EntityId = required_field(req.primary_target)?; + let s_target: EntityId = required_field(req.secondary_target)?; let (group, p_target, s_target) = app .write_tx(move |tx| { diff --git a/mgmtd/src/grpc/create_pool.rs b/mgmtd/src/grpc/create_pool.rs index 440c658f..2990ea44 100644 --- a/mgmtd/src/grpc/create_pool.rs +++ b/mgmtd/src/grpc/create_pool.rs @@ -14,8 +14,8 @@ pub(crate) async fn create_pool( bail!("node type must be storage"); } - let alias: Alias = required_field(req.alias)?.try_into()?; - let num_id: PoolId = req.num_id.unwrap_or_default().try_into()?; + let alias: Alias = required_field(req.alias)?; + let num_id: PoolId = optional_field(req.num_id)?.unwrap_or_default(); let (pool_uid, alias, pool_id) = app .write_tx(move |tx| { diff --git a/mgmtd/src/grpc/delete_buddy_group.rs b/mgmtd/src/grpc/delete_buddy_group.rs index 48d0c9f1..2c5cf07f 100644 --- a/mgmtd/src/grpc/delete_buddy_group.rs +++ b/mgmtd/src/grpc/delete_buddy_group.rs @@ -12,7 +12,7 @@ pub(crate) async fn delete_buddy_group( fail_on_missing_license(app, LicensedFeature::Mirroring)?; fail_on_pre_shutdown(app)?; - let group: EntityId = required_field(req.group)?.try_into()?; + let group: EntityId = required_field(req.group)?; let execute: bool = required_field(req.execute)?; // 1. Check deletion is allowed diff --git a/mgmtd/src/grpc/delete_node.rs b/mgmtd/src/grpc/delete_node.rs index 156ba947..88daa6af 100644 --- a/mgmtd/src/grpc/delete_node.rs +++ b/mgmtd/src/grpc/delete_node.rs @@ -8,7 +8,7 @@ pub(crate) async fn delete_node( ) -> Result { fail_on_pre_shutdown(app)?; - let node: EntityId = required_field(req.node)?.try_into()?; + let node: EntityId = required_field(req.node)?; let execute: bool = required_field(req.execute)?; let node = app diff --git a/mgmtd/src/grpc/delete_pool.rs b/mgmtd/src/grpc/delete_pool.rs index 23b5e988..0a908482 100644 --- a/mgmtd/src/grpc/delete_pool.rs +++ b/mgmtd/src/grpc/delete_pool.rs @@ -9,7 +9,7 @@ pub(crate) async fn delete_pool( fail_on_missing_license(app, LicensedFeature::Storagepool)?; fail_on_pre_shutdown(app)?; - let pool: EntityId = required_field(req.pool)?.try_into()?; + let pool: EntityId = required_field(req.pool)?; let execute: bool = required_field(req.execute)?; let pool = app diff --git a/mgmtd/src/grpc/delete_target.rs b/mgmtd/src/grpc/delete_target.rs index 1439949f..47b0b00e 100644 --- a/mgmtd/src/grpc/delete_target.rs +++ b/mgmtd/src/grpc/delete_target.rs @@ -9,7 +9,7 @@ pub(crate) async fn delete_target( ) -> Result { fail_on_pre_shutdown(app)?; - let target: EntityId = required_field(req.target)?.try_into()?; + let target: EntityId = required_field(req.target)?; let execute: bool = required_field(req.execute)?; let target = app diff --git a/mgmtd/src/grpc/set_alias.rs b/mgmtd/src/grpc/set_alias.rs index 800194e6..415f917a 100644 --- a/mgmtd/src/grpc/set_alias.rs +++ b/mgmtd/src/grpc/set_alias.rs @@ -11,7 +11,7 @@ pub(crate) async fn set_alias( // Parse proto msg let entity_type: EntityType = req.entity_type().try_into()?; - let entity_id: EntityId = required_field(req.entity_id)?.try_into()?; + let entity_id: EntityId = required_field(req.entity_id)?; let new_alias: Alias = req.new_alias.try_into()?; let update_alias_fn = move |tx: &Transaction, new_alias: &Alias| -> Result { diff --git a/mgmtd/src/grpc/set_default_quota_limits.rs b/mgmtd/src/grpc/set_default_quota_limits.rs index 127e38c6..e77b8c45 100644 --- a/mgmtd/src/grpc/set_default_quota_limits.rs +++ b/mgmtd/src/grpc/set_default_quota_limits.rs @@ -13,7 +13,7 @@ pub(crate) async fn set_default_quota_limits( bail!(QUOTA_NOT_ENABLED_STR); } - let pool: EntityId = required_field(req.pool)?.try_into()?; + let pool: EntityId = required_field(req.pool)?; fn update( tx: &Transaction, diff --git a/mgmtd/src/grpc/set_quota_limits.rs b/mgmtd/src/grpc/set_quota_limits.rs index e7f1b743..b8465d33 100644 --- a/mgmtd/src/grpc/set_quota_limits.rs +++ b/mgmtd/src/grpc/set_quota_limits.rs @@ -26,9 +26,9 @@ pub(crate) async fn set_quota_limits( for lim in req.limits { let id_type: QuotaIdType = lim.id_type().try_into()?; - let quota_id = required_field(lim.quota_id)?; + let quota_id: QuotaId = required_field(lim.quota_id)?; - let pool: EntityId = required_field(lim.pool)?.try_into()?; + let pool: EntityId = required_field(lim.pool)?; let pool_id = pool.resolve(tx, EntityType::Pool)?.num_id(); if let Some(l) = lim.space_limit { diff --git a/mgmtd/src/grpc/set_target_state.rs b/mgmtd/src/grpc/set_target_state.rs index 5eccc39a..b9b534d4 100644 --- a/mgmtd/src/grpc/set_target_state.rs +++ b/mgmtd/src/grpc/set_target_state.rs @@ -12,7 +12,7 @@ pub(crate) async fn set_target_state( fail_on_pre_shutdown(app)?; let state: TargetConsistencyState = req.consistency_state().try_into()?; - let target: EntityId = required_field(req.target)?.try_into()?; + let target: EntityId = required_field(req.target)?; let (target, node_uid) = app .write_tx(move |tx| { diff --git a/mgmtd/src/grpc/start_resync.rs b/mgmtd/src/grpc/start_resync.rs index 5a37c69a..ce70f88d 100644 --- a/mgmtd/src/grpc/start_resync.rs +++ b/mgmtd/src/grpc/start_resync.rs @@ -16,7 +16,7 @@ pub(crate) async fn start_resync( fail_on_missing_license(app, LicensedFeature::Mirroring)?; fail_on_pre_shutdown(app)?; - let buddy_group: EntityId = required_field(req.buddy_group)?.try_into()?; + let buddy_group: EntityId = required_field(req.buddy_group)?; let timestamp: i64 = required_field(req.timestamp)?; let restart: bool = required_field(req.restart)?; diff --git a/shared/src/grpc.rs b/shared/src/grpc.rs index 72365a9c..65b94db4 100644 --- a/shared/src/grpc.rs +++ b/shared/src/grpc.rs @@ -1,5 +1,5 @@ use anyhow::Result; -use std::fmt::Write; +use std::fmt::{Display, Write}; use std::future::Future; use std::pin::Pin; use tokio::sync::mpsc; @@ -216,10 +216,37 @@ pub fn process_grpc_handler_error(err: anyhow::Error) -> Status { Status::new(resp_code, err_string) } -/// Unwraps an optional proto message field . If `None`, errors out providing the fields name in the -/// error message. +/// Unwraps an optional protobuf message field while converting it to a local output type. If +/// `None`, errors out providing the fields name in the error message. /// /// Meant for unwrapping optional protobuf fields that are actually mandatory. -pub fn required_field(f: Option) -> Result { - f.ok_or_else(|| ::anyhow::anyhow!("missing required {} field", std::any::type_name::())) +pub fn required_field(f: Option) -> Result +where + R: TryFrom, + >::Error: Display, +{ + optional_field(f)? + .ok_or_else(|| anyhow::anyhow!("missing required {} field", std::any::type_name::())) +} + +/// Converts an optional protobuf message field into a local output type. +/// +/// Meant for unwrapping optional protobuf field that are actually optional (as defined in the +/// fields definition comment). If T is an enum, an error is thrown on the `unspecified` variant +/// which should never be set. +pub fn optional_field(f: Option) -> Result> +where + R: TryFrom, + >::Error: Display, +{ + f.map(|v| { + v.try_into().map_err(|err| { + anyhow::anyhow!( + "conversion of protobuf value of type {} to {} failed: {err:#}", + std::any::type_name::(), + std::any::type_name::(), + ) + }) + }) + .transpose() } diff --git a/shared/src/impl_macros.rs b/shared/src/impl_macros.rs index 70c5c0d1..d7004737 100644 --- a/shared/src/impl_macros.rs +++ b/shared/src/impl_macros.rs @@ -59,6 +59,7 @@ macro_rules! impl_enum_user_str { } #[cfg(feature = "grpc")] +#[macro_export] macro_rules! impl_enum_protobuf_traits { ($type:ty => $proto_type:ty, unspecified => $proto_unspec_variant:path, $($variant:path => $proto_variant:path),+ $(,)?) => { impl TryFrom<$proto_type> for $type { @@ -66,7 +67,7 @@ macro_rules! impl_enum_protobuf_traits { fn try_from(value: $proto_type) -> std::result::Result { let nt = match value { - $proto_unspec_variant => ::anyhow::bail!("$type is unspecified"), + $proto_unspec_variant => ::anyhow::bail!("{} is unspecified", stringify!($proto_type)), $( $proto_variant => $variant, )+ @@ -76,6 +77,16 @@ macro_rules! impl_enum_protobuf_traits { } } + impl TryFrom for $type { + type Error = ::anyhow::Error; + + fn try_from(value: i32) -> std::result::Result { + <$proto_type>::try_from(value) + .map_err(|_| {::anyhow::anyhow!( "{value} is not a valid {} value", stringify!($proto_type))})? + .try_into() + } + } + impl From<$type> for $proto_type { fn from(value: $type) -> Self { match value { From 408826b5dd095da9fe711ebd4e94e8a7258c5c92 Mon Sep 17 00:00:00 2001 From: Rusty Bee <145002912+rustybee42@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:39:32 +0200 Subject: [PATCH 2/9] fix: Improve quota and other logging * Quota: Log information after the operations are completed with more information at info level, remove the not so useful initial debug logs * Promote and improve some other state logging --- .../change_target_consistency_states.rs | 4 +- mgmtd/src/bee_msg/set_storage_target_info.rs | 4 +- mgmtd/src/quota.rs | 48 ++++++++++++------- 3 files changed, 37 insertions(+), 19 deletions(-) diff --git a/mgmtd/src/bee_msg/change_target_consistency_states.rs b/mgmtd/src/bee_msg/change_target_consistency_states.rs index 34c40513..6c449a4f 100644 --- a/mgmtd/src/bee_msg/change_target_consistency_states.rs +++ b/mgmtd/src/bee_msg/change_target_consistency_states.rs @@ -65,8 +65,8 @@ doesn't match stored state {old_stored}, no consistency state changes will be ma }) .await?; - log::debug!( - "Updated target states for {:?} targets {:?}, {} consistency states and {reachabilities_changed} reachability states changed", + log::info!( + "Updated {:?} targets' ({:?}) states: {} consistency states and {reachabilities_changed} reachability states changed", self.node_type, self.target_ids, consistencies_changed.unwrap_or(0) diff --git a/mgmtd/src/bee_msg/set_storage_target_info.rs b/mgmtd/src/bee_msg/set_storage_target_info.rs index 67646fe3..74714c90 100644 --- a/mgmtd/src/bee_msg/set_storage_target_info.rs +++ b/mgmtd/src/bee_msg/set_storage_target_info.rs @@ -15,6 +15,8 @@ impl HandleWithResponse for SetStorageTargetInfo { fail_on_pre_shutdown(app)?; let node_type = self.node_type; + let target_ids: Vec<_> = self.info.iter().map(|e| e.target_id).collect(); + app.write_tx(move |tx| { db::target::get_and_update_capacities( tx, @@ -34,7 +36,7 @@ impl HandleWithResponse for SetStorageTargetInfo { }) .await?; - log::debug!("Updated {node_type:?} target info"); + log::info!("Updated {node_type:?} targets' ({target_ids:?}) info and capacities"); // in the old mgmtd, a notice to refresh cap pools is sent out here if a cap pool // changed I consider this being to expensive to check here and just don't diff --git a/mgmtd/src/quota.rs b/mgmtd/src/quota.rs index f4422324..1a1d19fb 100644 --- a/mgmtd/src/quota.rs +++ b/mgmtd/src/quota.rs @@ -16,7 +16,9 @@ use sqlite::TransactionExt; use sqlite_check::sql; use std::collections::HashSet; use std::path::Path; +use std::sync::atomic::{AtomicUsize, Ordering}; use tokio::task::JoinHandle; +use tokio::time::Instant; #[derive(Debug, Clone, Copy)] struct TargetToQuery { @@ -59,10 +61,10 @@ pub(crate) async fn fetch_and_update(app: &impl App) -> Result<()> { return Ok(()); } - log::info!( - "Fetching quota information for {} storage targets", - targets_to_query.len() - ); + let targets_to_query_count = targets_to_query.len(); + + let start_time = Instant::now(); + let entry_counter = AtomicUsize::new(0); let tasks = create_and_send_requests(app, targets_to_query).await?; @@ -72,6 +74,8 @@ pub(crate) async fn fetch_and_update(app: &impl App) -> Result<()> { // Only process that target if there were not errors when fetching for this target if let Some(entries) = entries { + entry_counter.fetch_add(entries.len(), Ordering::Relaxed); + app.write_tx(move |tx| { // Always delete all the old entries for that target to make sure entries for no // longer queried ids are removed. We always get the complete list from the @@ -87,12 +91,6 @@ pub(crate) async fn fetch_and_update(app: &impl App) -> Result<()> { VALUES (?1, ?2, ?3 ,?4 ,?5)" ))?; - log::debug!( - "Setting {} quota usage entries for target {}", - entries.len(), - target.target_id - ); - for e in entries { if e.space > 0 { insert_stmt.execute(params![ @@ -121,6 +119,13 @@ pub(crate) async fn fetch_and_update(app: &impl App) -> Result<()> { } } + log::info!( + "Fetched and stored {} quota entries from {} targets in {:?}", + entry_counter.load(Ordering::Relaxed), + targets_to_query_count, + start_time.elapsed() + ); + Ok(()) } @@ -326,7 +331,6 @@ pub(crate) async fn distribute_exceeded(app: &impl App) -> Result<()> { if !app.static_info().user_config.quota_enforce { return Ok(()); } - log::info!("Calculating and pushing exceeded quota"); let quota_licensed = app.verify_licensed_feature(LicensedFeature::Quota).is_ok(); @@ -376,6 +380,8 @@ pub(crate) async fn distribute_exceeded(app: &impl App) -> Result<()> { } } } else { + // If quota is unlicensed, make sure the exceeding ids are removed from the servers. + // Otherwise exceeded ids could stay exceeded forever if quota was used before. log::info!( "Quota enforcement enabled but feature not licensed. Removing quota limits from nodes" ); @@ -396,20 +402,22 @@ pub(crate) async fn distribute_exceeded(app: &impl App) -> Result<()> { }) .await?; + let start_time = Instant::now(); + let mut id_counter = 0; + // Send all messages with exceeded quota information to all meta and storage nodes // Since there is one message for each combination of (pool x (user, group) x (space, inode)), // this might be very demanding, but can't do anything about that without changing meta and // storage too. // If this shows as a bottleneck, the requests could be done concurrently though. - for msg in msges { + for msg in &msges { let mut request_fails = 0; let mut non_success_count = 0; + id_counter += msg.exceeded_quota_ids.len(); + for node_uid in &nodes { - match app - .request::<_, SetExceededQuotaResp>(*node_uid, &msg) - .await - { + match app.request::<_, SetExceededQuotaResp>(*node_uid, msg).await { Ok(resp) => { if resp.result != OpsErr::SUCCESS { non_success_count += 1; @@ -429,6 +437,14 @@ pub(crate) async fn distribute_exceeded(app: &impl App) -> Result<()> { } } + log::info!( + "Pushed {} exceeded quota ids to {} nodes using {} messages in {:?}", + id_counter, + nodes.len(), + msges.len(), + start_time.elapsed() + ); + Ok(()) } From f2bd21c19e70d9ec34af3c44460a9667c8b681a3 Mon Sep 17 00:00:00 2001 From: Rusty Bee <145002912+rustybee42@users.noreply.github.com> Date: Wed, 12 Aug 2026 14:19:05 +0200 Subject: [PATCH 3/9] chore: Update rusqlite Required to get newest the newest sqlite, which allows adding constraints using ALTER TABLE --- Cargo.lock | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8f15e3a3..aa746ec6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -417,14 +417,17 @@ name = "hashbrown" version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4f467dd6dccf739c208452f8014c75c18bb8301b050ad1cfb27153803edb0f51" +dependencies = [ + "foldhash 0.2.0", +] [[package]] name = "hashlink" -version = "0.11.0" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea0b22561a9c04a7cb1a302c013e0259cd3b4bb619f145b32f72b8b4bcbed230" +checksum = "32069d97bb81e38fa67eab65e3393bf804bb85969f2bc06bf13f64aef5aba248" dependencies = [ - "hashbrown 0.16.1", + "hashbrown 0.17.0", ] [[package]] @@ -639,9 +642,9 @@ dependencies = [ [[package]] name = "libsqlite3-sys" -version = "0.37.0" +version = "0.38.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1f111c8c41e7c61a49cd34e44c7619462967221a6443b0ec299e0ac30cfb9b1" +checksum = "f1d20bef17f513b9b3004532233187769cd072d790971f4e4da0e346eb6401e8" dependencies = [ "cc", "pkg-config", @@ -947,9 +950,9 @@ dependencies = [ [[package]] name = "rusqlite" -version = "0.39.0" +version = "0.40.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a0d2b0146dd9661bf67bb107c0bb2a55064d556eeb3fc314151b957f313bcd4e" +checksum = "23f2a97da3e3873c73cb2a2e71b35c40ff95e0b1eefa8d72d8499a6928c3b5b3" dependencies = [ "bitflags", "fallible-iterator", From b44880da2e8e84597a9ea93ef21f21855b128799 Mon Sep 17 00:00:00 2001 From: Rusty Bee <145002912+rustybee42@users.noreply.github.com> Date: Wed, 12 Aug 2026 11:59:32 +0200 Subject: [PATCH 4/9] feat: Exclude secondary targets from quota accounting by default Before, targets in buddy groups would count double towards the limits as quota would just use the data from both. Now, only the primary target is taken into account by default. This has the consequence that potentially unmirrored data present on the secondary will not be counted towards the quota anymore. And, if there is a switchover, the numbers might suddenly change. We assume though, that most customers have exclusive mirrored targets. To get the old behavior, the accounting mode is stored with the buddy groups and can be set via ctl. On updating management (from 8.x), it will use the old behavior, adding new groups uses the new behavior. Upgrading from v7 uses the new behavior as this usually makes more sense and is already a breaking change. * Also add modify_buddy_group() handler to change the accounting mode and adapt the create_buddy_group() and get_buddy_groups() handlers * Refactor the get_buddy_groups() big sql query, using string for selecting columns --- Cargo.lock | 2 +- Cargo.toml | 2 +- mgmtd/src/db/buddy_group.rs | 22 +++++++++++-- mgmtd/src/db/import_v7.rs | 1 + mgmtd/src/db/schema/7.sql | 7 +++++ mgmtd/src/db/schema/test_data.sql | 8 ++--- mgmtd/src/grpc.rs | 6 ++++ mgmtd/src/grpc/create_buddy_group.rs | 14 +++++++++ mgmtd/src/grpc/get_buddy_groups.rs | 46 ++++++++++++++++------------ mgmtd/src/grpc/get_quota_usage.rs | 17 ++++++++-- mgmtd/src/grpc/modify_buddy_group.rs | 41 +++++++++++++++++++++++++ mgmtd/src/quota.rs | 7 +++-- mgmtd/src/types.rs | 25 +++++++++++++-- 13 files changed, 163 insertions(+), 35 deletions(-) create mode 100644 mgmtd/src/db/schema/7.sql create mode 100644 mgmtd/src/grpc/modify_buddy_group.rs diff --git a/Cargo.lock b/Cargo.lock index aa746ec6..c38dfabb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -872,7 +872,7 @@ dependencies = [ [[package]] name = "protobuf" version = "0.0.0" -source = "git+https://github.com/thinkparq/protobuf?rev=4d5e5db085065acbbaa5bb76ce4b81d6d733e446#4d5e5db085065acbbaa5bb76ce4b81d6d733e446" +source = "git+https://github.com/thinkparq/protobuf?rev=25d79b6293ee46e1fddf832874836472951aeb29#25d79b6293ee46e1fddf832874836472951aeb29" dependencies = [ "prost", "prost-types", diff --git a/Cargo.toml b/Cargo.toml index d15d8155..e4d42523 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,7 +17,7 @@ itertools = "0" libc = "0" log = { version = "0", features = ["std"] } prost = "0.14" -protobuf = { git = "https://github.com/thinkparq/protobuf", rev = "4d5e5db085065acbbaa5bb76ce4b81d6d733e446" } +protobuf = { git = "https://github.com/thinkparq/protobuf", rev = "25d79b6293ee46e1fddf832874836472951aeb29" } regex = "1" ring = "0" rusqlite = { version = "0", features = ["bundled", "vtab", "array", "fallible_uint"] } diff --git a/mgmtd/src/db/buddy_group.rs b/mgmtd/src/db/buddy_group.rs index 32085bb2..52a78b1c 100644 --- a/mgmtd/src/db/buddy_group.rs +++ b/mgmtd/src/db/buddy_group.rs @@ -48,6 +48,7 @@ pub(crate) fn insert( node_type: NodeTypeServer, p_target_id: TargetId, s_target_id: TargetId, + quota_accounting: Option, ) -> Result<(Uid, BuddyGroupId)> { let group_id = if group_id == 0 { misc::find_new_id(tx, "buddy_groups", "group_id", node_type.into(), 1..=0xFFFF)? @@ -128,12 +129,24 @@ pub(crate) fn insert( None }; + // Quota accounting only applies to storage groups + let quota_accounting = match (node_type, quota_accounting) { + (NodeTypeServer::Meta, None) => None, + (NodeTypeServer::Meta, Some(_)) => { + bail!("The quota accounting mode can only be set for storage buddy groups") + } + (NodeTypeServer::Storage, None) => { + bail!("The quota accounting mode must be set for storage buddy groups"); + } + (NodeTypeServer::Storage, Some(q)) => Some(q), + }; + // Insert generic buddy group tx.execute( sql!( "INSERT INTO buddy_groups - (group_uid, node_type, group_id, p_target_id, s_target_id, pool_id) - VALUES (?1, ?2, ?3, ?4, ?5, ?6)" + (group_uid, node_type, group_id, p_target_id, s_target_id, pool_id, quota_accounting) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)" ), params![ new_uid, @@ -141,7 +154,8 @@ pub(crate) fn insert( group_id, p_target_id, s_target_id, - pool_id + pool_id, + quota_accounting.map(|e| e.sql_variant()) ], )?; @@ -298,6 +312,7 @@ mod test { NodeTypeServer::Meta, 3, 4, + None, ) .unwrap(); super::insert( @@ -307,6 +322,7 @@ mod test { NodeTypeServer::Storage, 3, 7, + None, ) .unwrap_err(); diff --git a/mgmtd/src/db/import_v7.rs b/mgmtd/src/db/import_v7.rs index d115da9b..663a8da7 100644 --- a/mgmtd/src/db/import_v7.rs +++ b/mgmtd/src/db/import_v7.rs @@ -255,6 +255,7 @@ fn buddy_groups(tx: &Transaction, f: &Path, nt: NodeTypeServer) -> Result<()> { nt, BuddyGroupId::from_str_radix(p_id.trim(), 16)?, BuddyGroupId::from_str_radix(s_id.trim(), 16)?, + matches!(nt, NodeTypeServer::Storage).then_some(BuddyGroupQuotaAccounting::Both), )?; } diff --git a/mgmtd/src/db/schema/7.sql b/mgmtd/src/db/schema/7.sql new file mode 100644 index 00000000..3bc01d87 --- /dev/null +++ b/mgmtd/src/db/schema/7.sql @@ -0,0 +1,7 @@ +-- This can be NULL as it only applies to storage groups +ALTER table buddy_groups ADD COLUMN quota_accounting INTEGER; +-- Make sure existing mirrored targets don't change behavior automatically +UPDATE buddy_groups SET quota_accounting = 2 WHERE node_type = 2; + +ALTER table buddy_groups ADD CONSTRAINT quota_accounting_null +CHECK ((node_type == 2) == (quota_accounting IS NOT NULL)); diff --git a/mgmtd/src/db/schema/test_data.sql b/mgmtd/src/db/schema/test_data.sql index b1d8564e..13363ca3 100644 --- a/mgmtd/src/db/schema/test_data.sql +++ b/mgmtd/src/db/schema/test_data.sql @@ -137,10 +137,10 @@ INSERT INTO entities (uid, entity_type, alias) VALUES (302002, 4, "storage_buddy_group_2") ; -INSERT INTO buddy_groups (group_uid, node_type, group_id, p_target_id, s_target_id, pool_id) VALUES - (301001, 1, 1, 1, 2, NULL), - (302001, 2, 1, 1, 5, 1), - (302002, 2, 2, 9, 13, 1) +INSERT INTO buddy_groups (group_uid, node_type, group_id, p_target_id, s_target_id, pool_id, quota_accounting) VALUES + (301001, 1, 1, 1, 2, NULL, NULL), + (302001, 2, 1, 1, 5, 1, 1), + (302002, 2, 2, 9, 13, 1, 1) ; diff --git a/mgmtd/src/grpc.rs b/mgmtd/src/grpc.rs index 8292d6ff..220d99d7 100644 --- a/mgmtd/src/grpc.rs +++ b/mgmtd/src/grpc.rs @@ -37,6 +37,7 @@ mod get_quota_limits; mod get_quota_usage; mod get_targets; mod mirror_root_inode; +mod modify_buddy_group; mod set_alias; mod set_default_quota_limits; mod set_quota_limits; @@ -126,6 +127,11 @@ impl pm::management_server::Management for ManagementService { pm::CreateBuddyGroupRequest => pm::CreateBuddyGroupResponse, "Create buddy group" } + impl_grpc_handler! { + modify_buddy_group, + pm::ModifyBuddyGroupRequest => pm::ModifyBuddyGroupResponse, + "Modify buddy group" + } impl_grpc_handler! { delete_buddy_group, pm::DeleteBuddyGroupRequest => pm::DeleteBuddyGroupResponse, diff --git a/mgmtd/src/grpc/create_buddy_group.rs b/mgmtd/src/grpc/create_buddy_group.rs index a7cf97e4..0d71aff4 100644 --- a/mgmtd/src/grpc/create_buddy_group.rs +++ b/mgmtd/src/grpc/create_buddy_group.rs @@ -1,4 +1,5 @@ use super::*; +use crate::types::BuddyGroupQuotaAccounting; use shared::bee_msg::buddy_group::SetMirrorBuddyGroup; use shared::bee_msg::storage_pool::RefreshStoragePools; @@ -16,6 +17,18 @@ pub(crate) async fn create_buddy_group( let p_target: EntityId = required_field(req.primary_target)?; let s_target: EntityId = required_field(req.secondary_target)?; + // Compatibility: The options field is optional. + let options: pm::BuddyGroupOptions = optional_field(req.options)?.unwrap_or_default(); + let quota_accounting: Option = + optional_field(options.quota_accounting)?.or_else(|| { + // Compatibility: Allow creating storage buddy groups without this field by defaulting. + if node_type == NodeTypeServer::Storage { + Some(BuddyGroupQuotaAccounting::Primary) + } else { + None + } + }); + let (group, p_target, s_target) = app .write_tx(move |tx| { let p_target = p_target.resolve(tx, EntityType::Target)?; @@ -28,6 +41,7 @@ pub(crate) async fn create_buddy_group( node_type, p_target.num_id().try_into()?, s_target.num_id().try_into()?, + quota_accounting, )?; Ok(( EntityIdSet { diff --git a/mgmtd/src/grpc/get_buddy_groups.rs b/mgmtd/src/grpc/get_buddy_groups.rs index fc8fec0c..02e692df 100644 --- a/mgmtd/src/grpc/get_buddy_groups.rs +++ b/mgmtd/src/grpc/get_buddy_groups.rs @@ -9,11 +9,12 @@ pub(crate) async fn get_buddy_groups( .read_tx(|tx| { Ok(tx.query_map_collect( sql!( - "SELECT group_uid, group_id, bg.alias, bg.node_type, - p_target_uid, p_t.target_id, p_t.alias, - s_target_uid, s_t.target_id, s_t.alias, - p.pool_uid, bg.pool_id, p.alias, - p_t.consistency, s_t.consistency + "SELECT group_uid, group_id, bg.alias AS group_alias, bg.node_type, + p_target_uid, p_t.target_id AS p_target_id, p_t.alias AS p_target_alias, + s_target_uid, s_t.target_id AS s_target_id, s_t.alias AS s_target_alias, + p.pool_uid, bg.pool_id, p.alias AS pool_alias, + p_t.consistency AS p_consistency, s_t.consistency AS s_consistency, + bg.quota_accounting FROM buddy_groups_ext AS bg INNER JOIN targets_ext AS p_t ON p_t.target_uid = p_target_uid INNER JOIN targets_ext AS s_t ON s_t.target_uid = s_target_uid @@ -21,50 +22,55 @@ pub(crate) async fn get_buddy_groups( ), [], |row| { - let node_type = NodeType::from_row(row, 3)?.into_proto_i32(); - let p_con_state = TargetConsistencyState::from_row(row, 13)?.into_proto_i32(); - let s_con_state = TargetConsistencyState::from_row(row, 14)?.into_proto_i32(); + let node_type = NodeType::from_row(row, "node_type")?.into_proto_i32(); + let p_con_state = + TargetConsistencyState::from_row(row, "p_consistency")?.into_proto_i32(); + let s_con_state = + TargetConsistencyState::from_row(row, "s_consistency")?.into_proto_i32(); Ok(pm::get_buddy_groups_response::BuddyGroup { id: Some(pb::EntityIdSet { - uid: row.get(0)?, + uid: row.get("group_uid")?, legacy_id: Some(pb::LegacyId { - num_id: row.get(1)?, + num_id: row.get("group_id")?, node_type, }), - alias: row.get(2)?, + alias: row.get("group_alias")?, }), node_type, primary_target: Some(pb::EntityIdSet { - uid: row.get(4)?, + uid: row.get("p_target_uid")?, legacy_id: Some(pb::LegacyId { - num_id: row.get(5)?, + num_id: row.get("p_target_id")?, node_type, }), - alias: row.get(6)?, + alias: row.get("p_target_alias")?, }), secondary_target: Some(pb::EntityIdSet { - uid: row.get(7)?, + uid: row.get("s_target_uid")?, legacy_id: Some(pb::LegacyId { - num_id: row.get(8)?, + num_id: row.get("s_target_id")?, node_type, }), - alias: row.get(9)?, + alias: row.get("s_target_alias")?, }), - storage_pool: if let Some(uid) = row.get::<_, Option>(10)? { + storage_pool: if let Some(uid) = row.get::<_, Option>("pool_uid")? { Some(pb::EntityIdSet { uid: Some(uid), legacy_id: Some(pb::LegacyId { - num_id: row.get(11)?, + num_id: row.get("pool_id")?, node_type, }), - alias: row.get(12)?, + alias: row.get("pool_alias")?, }) } else { None }, primary_consistency_state: p_con_state, secondary_consistency_state: s_con_state, + options: Some(pm::BuddyGroupOptions { + quota_accounting: row.get("quota_accounting")?, + }), }) }, )?) diff --git a/mgmtd/src/grpc/get_quota_usage.rs b/mgmtd/src/grpc/get_quota_usage.rs index cd6ed5f2..82598776 100644 --- a/mgmtd/src/grpc/get_quota_usage.rs +++ b/mgmtd/src/grpc/get_quota_usage.rs @@ -1,5 +1,6 @@ use super::common::{QUOTA_NOT_ENABLED_STR, QUOTA_STREAM_BUF_SIZE, QUOTA_STREAM_PAGE_LIMIT}; use super::*; +use crate::types::BuddyGroupQuotaAccounting; use itertools::Itertools; use std::fmt::Write; @@ -13,8 +14,9 @@ pub(crate) async fn get_quota_usage( bail!(QUOTA_NOT_ENABLED_STR); } - let mut r#where = "FALSE ".to_string(); + let mut r#where = "(FALSE ".to_string(); + // Optionally filter by id range or list, separate for uids and gids let mut filter = |min: Option, max: Option, list: &[u32], typ: QuotaIdType| -> Result<()> { if min.is_some() || max.is_some() || !list.is_empty() { @@ -52,6 +54,15 @@ pub(crate) async fn get_quota_usage( QuotaIdType::Group, )?; + write!(r#where, ") ")?; + + // Filter out secondary targets if configured to avoid double accounting on buddy groups + write!( + r#where, + "AND (bg.quota_accounting IS NULL OR bg.quota_accounting = {})", + BuddyGroupQuotaAccounting::Both.sql_variant() + )?; + let mut having = "TRUE ".to_string(); if let Some(pool) = req.pool { @@ -65,7 +76,7 @@ pub(crate) async fn get_quota_usage( } if let Some(exceeded) = req.exceeded { let base = "(space_used > space_limit AND space_limit > -1 - OR inode_used > inode_limit AND inode_limit > -1)"; + OR inode_used > inode_limit AND inode_limit > -1)"; if exceeded { write!(having, "AND {base} ")?; } else { @@ -85,6 +96,8 @@ pub(crate) async fn get_quota_usage( SUM(CASE WHEN u.quota_type = {inode} THEN u.value END) AS inode_used FROM quota_usage AS u INNER JOIN targets AS st USING(node_type, target_id) + LEFT JOIN buddy_groups AS bg ON st.target_id = bg.s_target_id + AND st.node_type = bg.node_type INNER JOIN pools_ext AS sp USING(node_type, pool_id) LEFT JOIN quota_default_limits AS d USING(id_type, quota_type, pool_id) LEFT JOIN quota_limits AS l USING(quota_id, id_type, quota_type, pool_id) diff --git a/mgmtd/src/grpc/modify_buddy_group.rs b/mgmtd/src/grpc/modify_buddy_group.rs new file mode 100644 index 00000000..bc5762f6 --- /dev/null +++ b/mgmtd/src/grpc/modify_buddy_group.rs @@ -0,0 +1,41 @@ +use super::*; +use crate::types::BuddyGroupQuotaAccounting; + +/// Modify the settings of an existing buddy group +pub(crate) async fn modify_buddy_group( + app: &impl App, + req: pm::ModifyBuddyGroupRequest, +) -> Result { + fail_on_missing_license(app, LicensedFeature::Mirroring)?; + fail_on_pre_shutdown(app)?; + + let group: EntityId = required_field(req.group)?; + let options: pm::BuddyGroupOptions = required_field(req.options)?; + let quota_accounting: Option = + optional_field(options.quota_accounting)?; + + let group = app + .write_tx(move |tx| { + let group = group.resolve(tx, EntityType::BuddyGroup)?; + + if quota_accounting.is_some() && group.node_type() != NodeType::Storage { + bail!("The quota accounting mode can only be set to storage buddy groups"); + } + + tx.execute_cached( + sql!( + "UPDATE buddy_groups + SET quota_accounting = COALESCE(?1, quota_accounting) + WHERE group_uid = ?2" + ), + params![quota_accounting.map(|e| e.sql_variant()), group.uid], + )?; + + Ok(group) + }) + .await?; + + log::info!("Buddy group {group} modified: quota_accounting={quota_accounting:?}"); + + Ok(pm::ModifyBuddyGroupResponse {}) +} diff --git a/mgmtd/src/quota.rs b/mgmtd/src/quota.rs index 1a1d19fb..900e536d 100644 --- a/mgmtd/src/quota.rs +++ b/mgmtd/src/quota.rs @@ -4,7 +4,7 @@ mod system_id; use crate::app::*; use crate::license::LicensedFeature; -use crate::types::SqliteEnumExt; +use crate::types::{BuddyGroupQuotaAccounting, SqliteEnumExt}; use anyhow::{Context as AnyhowContext, Result}; use rusqlite::params; use shared::bee_msg::OpsErr; @@ -362,12 +362,15 @@ pub(crate) async fn distribute_exceeded(app: &impl App) -> Result<()> { "SELECT DISTINCT e.quota_id, e.id_type, e.quota_type, st.pool_id FROM quota_usage AS e INNER JOIN targets AS st USING(node_type, target_id) + LEFT JOIN buddy_groups AS bg ON st.target_id = bg.s_target_id + AND st.node_type = bg.node_type LEFT JOIN quota_default_limits AS d USING(id_type, quota_type, pool_id) LEFT JOIN quota_limits AS l USING(quota_id, id_type, quota_type, pool_id) + WHERE bg.quota_accounting IS NULL OR bg.quota_accounting = ?1 GROUP BY e.quota_id, e.id_type, e.quota_type, st.pool_id HAVING SUM(e.value) > COALESCE(l.value, d.value)" ))?; - let mut rows = stmt.query([])?; + let mut rows = stmt.query([BuddyGroupQuotaAccounting::Both.sql_variant()])?; while let Some(row) = rows.next()? { for m in &mut msges { if row.get::<_, PoolId>(3)? == m.pool_id diff --git a/mgmtd/src/types.rs b/mgmtd/src/types.rs index 8edbe686..53331cf7 100644 --- a/mgmtd/src/types.rs +++ b/mgmtd/src/types.rs @@ -1,6 +1,8 @@ //! Contains types used by the local database and config. -use rusqlite::Row; +use protobuf::management as pm; +use rusqlite::{Row, RowIndex}; +use shared::impl_enum_protobuf_traits; use shared::types::*; mod entity; @@ -13,7 +15,7 @@ pub(crate) trait SqliteEnumExt { where Self: Sized; - fn from_row(row: &Row, idx: usize) -> rusqlite::Result + fn from_row(row: &Row, idx: impl RowIndex) -> rusqlite::Result where Self: Sized, { @@ -87,3 +89,22 @@ impl_enum_sqlite! {QuotaType, QuotaType::Space => 1, QuotaType::Inode => 2, } + +/// How to handle quota data on targets that are part of a buddy group +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub(crate) enum BuddyGroupQuotaAccounting { + Primary, + Both, +} + +impl_enum_sqlite! {BuddyGroupQuotaAccounting, + BuddyGroupQuotaAccounting::Primary => 1, + BuddyGroupQuotaAccounting::Both => 2, +} + +use pm::buddy_group_options::BuddyGroupQuotaAccounting as BGQ; +impl_enum_protobuf_traits! {BuddyGroupQuotaAccounting => BGQ, + unspecified => BGQ::Unspecified, + BuddyGroupQuotaAccounting::Primary => BGQ::Primary, + BuddyGroupQuotaAccounting::Both => BGQ::Both, +} From cd5861bfc88fc553640f419dc18b9b51c1c4b8a6 Mon Sep 17 00:00:00 2001 From: Rusty Bee <145002912+rustybee42@users.noreply.github.com> Date: Thu, 13 Aug 2026 14:32:25 +0200 Subject: [PATCH 5/9] refactor: improve quota tests, remove exceeded statement duplication Test the accounting modes, making sure they work as expected + more. --- mgmtd/src/bee_msg/request_exceeded_quota.rs | 15 ++--- mgmtd/src/db/buddy_group.rs | 2 +- mgmtd/src/db/misc.rs | 2 +- mgmtd/src/db/node.rs | 6 +- mgmtd/src/db/schema/test_data.sql | 23 ++++++- mgmtd/src/grpc/get_nodes.rs | 4 +- mgmtd/src/quota.rs | 74 +++++++++++++++------ 7 files changed, 86 insertions(+), 40 deletions(-) diff --git a/mgmtd/src/bee_msg/request_exceeded_quota.rs b/mgmtd/src/bee_msg/request_exceeded_quota.rs index 17ba72f9..e83b9b89 100644 --- a/mgmtd/src/bee_msg/request_exceeded_quota.rs +++ b/mgmtd/src/bee_msg/request_exceeded_quota.rs @@ -30,20 +30,13 @@ impl HandleWithResponse for RequestExceededQuota { )? }; + // Query the exceeded ids matching the request let exceeded_quota_ids = tx.query_map_collect( - sql!( - "SELECT DISTINCT e.quota_id FROM quota_usage AS e - INNER JOIN targets AS st USING(node_type, target_id) - LEFT JOIN quota_default_limits AS d USING(id_type, quota_type, pool_id) - LEFT JOIN quota_limits AS l USING(quota_id, id_type, quota_type, pool_id) - WHERE e.id_type = ?1 AND e.quota_type = ?2 AND st.pool_id = ?3 - GROUP BY e.quota_id, e.id_type, e.quota_type, st.pool_id - HAVING SUM(e.value) > COALESCE(l.value, d.value)" - ), + crate::quota::EXCEEDED_QUOTA_IDS_SQL, params![ self.id_type.sql_variant(), self.quota_type.sql_variant(), - pool_id + pool_id, ], |row| row.get(0), )?; @@ -84,7 +77,7 @@ mod test { pool_id: 1, target_id: 0, }, - &[2, 4, 10], + &[2, 4, 10, 51], ), ( RequestExceededQuota { diff --git a/mgmtd/src/db/buddy_group.rs b/mgmtd/src/db/buddy_group.rs index 52a78b1c..f6c51c34 100644 --- a/mgmtd/src/db/buddy_group.rs +++ b/mgmtd/src/db/buddy_group.rs @@ -329,7 +329,7 @@ mod test { let meta_groups = get_with_type(tx, NodeTypeServer::Meta).unwrap(); let storage_groups = get_with_type(tx, NodeTypeServer::Storage).unwrap(); - assert_eq!(2, meta_groups.len()); + assert_eq!(3, meta_groups.len()); assert_eq!(2, storage_groups.len()); assert!(meta_groups.iter().any(|e| e.0 == 1234)); }) diff --git a/mgmtd/src/db/misc.rs b/mgmtd/src/db/misc.rs index 259fcb78..aff7500f 100644 --- a/mgmtd/src/db/misc.rs +++ b/mgmtd/src/db/misc.rs @@ -132,7 +132,7 @@ mod test { // New max id let new_id = super::find_new_id(tx, "targets", "target_id", NodeType::Meta, 1..=100).unwrap(); - assert_eq!(new_id, 5); + assert_eq!(new_id, 6); // New min ID in a non-empty range let new_id = super::find_new_id(tx, "targets", "target_id", NodeType::Meta, 0..=4).unwrap(); diff --git a/mgmtd/src/db/node.rs b/mgmtd/src/db/node.rs index 0ecce5c0..7ddffb7e 100644 --- a/mgmtd/src/db/node.rs +++ b/mgmtd/src/db/node.rs @@ -208,7 +208,7 @@ mod test { #[test] fn insert_get_delete() { with_test_data(|tx| { - assert_eq!(5, get_with_type(tx, NodeType::Meta).unwrap().len()); + assert_eq!(7, get_with_type(tx, NodeType::Meta).unwrap().len()); let node = insert( tx, 1234, @@ -233,11 +233,11 @@ mod test { 10000, ) .unwrap_err(); - assert_eq!(6, get_with_type(tx, NodeType::Meta).unwrap().len()); + assert_eq!(8, get_with_type(tx, NodeType::Meta).unwrap().len()); delete(tx, node.uid).unwrap(); delete(tx, node.uid).unwrap_err(); - assert_eq!(5, get_with_type(tx, NodeType::Meta).unwrap().len()); + assert_eq!(7, get_with_type(tx, NodeType::Meta).unwrap().len()); }); } diff --git a/mgmtd/src/db/schema/test_data.sql b/mgmtd/src/db/schema/test_data.sql index 13363ca3..97989a99 100644 --- a/mgmtd/src/db/schema/test_data.sql +++ b/mgmtd/src/db/schema/test_data.sql @@ -11,6 +11,8 @@ INSERT INTO entities (uid, entity_type, alias) VALUES (101002, 1, "meta_node_2"), (101003, 1, "meta_node_3"), (101004, 1, "meta_node_4"), + (101005, 1, "meta_node_5"), + (101013, 1, "meta_node_13"), (101099, 1, "meta_node_no_target"), (102001, 1, "storage_node_1"), (102002, 1, "storage_node_2"), @@ -27,6 +29,8 @@ INSERT INTO nodes (node_uid, node_id, node_type, port, last_contact) VALUES (101002, 2, 1, 8005, DATETIME("NOW")), (101003, 3, 1, 8005, DATETIME("NOW")), (101004, 4, 1, 8005, DATETIME("NOW")), + (101005, 5, 1, 8005, DATETIME("NOW")), + (101013, 13, 1, 8005, DATETIME("NOW")), (101099, 99, 1, 8005, DATETIME("NOW")), @@ -83,6 +87,8 @@ INSERT INTO entities (uid, entity_type, alias) VALUES (201002, 2, "meta_target_2"), (201003, 2, "meta_target_3"), (201004, 2, "meta_target_4"), + (201005, 2, "meta_target_5"), + (201013, 2, "meta_target_13"), (202001, 2, "storage_target_1"), (202002, 2, "storage_target_2"), @@ -110,6 +116,8 @@ free_space, free_inodes, consistency, last_update) VALUES (201002, 1, 2, 2, NULL, 1000000, 1000000, 550000, 550000, 1, DATETIME("NOW")), (201003, 1, 3, 3, NULL, 1000000, 1000000, 550000, 550000, 1, DATETIME("NOW")), (201004, 1, 4, 4, NULL, 1000000, 1000000, 450000, 450000, 1, DATETIME("NOW")), + (201005, 1, 5, 3, NULL, 1000000, 1000000, 550000, 550000, 1, DATETIME("NOW")), + (201013, 1, 13, 4, NULL, 1000000, 1000000, 450000, 450000, 1, DATETIME("NOW")), (202001, 2, 1, 1, 1, 1000000, 1000000, 450000, 450000, 1, DATETIME("NOW")), (202002, 2, 2, 1, 2, 1000000, 1000000, 500000, 500000, 1, DATETIME("NOW")), @@ -133,14 +141,17 @@ free_space, free_inodes, consistency, last_update) VALUES INSERT INTO entities (uid, entity_type, alias) VALUES (301001, 4, "meta_buddy_group_1"), + (301002, 4, "meta_buddy_group_2"), (302001, 4, "storage_buddy_group_1"), (302002, 4, "storage_buddy_group_2") ; INSERT INTO buddy_groups (group_uid, node_type, group_id, p_target_id, s_target_id, pool_id, quota_accounting) VALUES (301001, 1, 1, 1, 2, NULL, NULL), + -- this meta buddy group tests target id separation from storage buddy groups + (301002, 1, 2, 5, 13, NULL, NULL), (302001, 2, 1, 1, 5, 1, 1), - (302002, 2, 2, 9, 13, 1, 1) + (302002, 2, 2, 9, 13, 1, 2) ; @@ -192,5 +203,13 @@ INSERT INTO quota_usage (quota_id, id_type, quota_type, target_id, value) VALUES (10, 2, 1, 2, 999999999), (10, 1, 2, 2, 999999999), (10, 2, 2, 2, 999999999), - (20, 1, 1, 2, 101) + (20, 1, 1, 2, 101), + -- 1 + 5 has quota accounting mode "primary", counted once, not exceeeded + (50, 1, 1, 1, 800), + (50, 1, 1, 5, 800), + -- 9 + 13 has quota accounting mode "both", counted twice, exceeded + (51, 1, 1, 9, 800), + (51, 1, 1, 13, 800), + -- target 13 exists both as meta and storage, this tests their correct separation - not exceeded + (52, 1, 1, 13, 800) ; diff --git a/mgmtd/src/grpc/get_nodes.rs b/mgmtd/src/grpc/get_nodes.rs index dac810b7..3e2a5547 100644 --- a/mgmtd/src/grpc/get_nodes.rs +++ b/mgmtd/src/grpc/get_nodes.rs @@ -177,14 +177,14 @@ mod test { .await .unwrap(); - assert_eq!(res.nodes.len(), 14); + assert_eq!(res.nodes.len(), 16); assert!(res.nodes.iter().all(|e| e.nics.is_empty())); let res = super::get_nodes(&app, pm::GetNodesRequest { include_nics: true }) .await .unwrap(); - assert_eq!(res.nodes.len(), 14); + assert_eq!(res.nodes.len(), 16); assert_eq!( res.nodes .iter() diff --git a/mgmtd/src/quota.rs b/mgmtd/src/quota.rs index 900e536d..f159c700 100644 --- a/mgmtd/src/quota.rs +++ b/mgmtd/src/quota.rs @@ -4,7 +4,7 @@ mod system_id; use crate::app::*; use crate::license::LicensedFeature; -use crate::types::{BuddyGroupQuotaAccounting, SqliteEnumExt}; +use crate::types::SqliteEnumExt; use anyhow::{Context as AnyhowContext, Result}; use rusqlite::params; use shared::bee_msg::OpsErr; @@ -326,6 +326,31 @@ fn extract_results( } } +/// Finds exceeded quota ids +/// +/// The three parameters can be set to filter the data put into the result (before grouping) or set +/// to `None` to get everything. This uses a hardcoded `2 = both` for the quota accounting mode - +/// usage on a secondary is only counted when set to that mode. +/// +/// Note that `quota_usage` is scanned either way: its primary key starts with `quota_id`, so +/// neither the fixed nor the optional form of the id type / quota type filters can seek on it (thus +/// no difference in performance). +pub(crate) const EXCEEDED_QUOTA_IDS_SQL: &str = sql!( + "SELECT DISTINCT e.quota_id, e.id_type, e.quota_type, st.pool_id + FROM quota_usage AS e + INNER JOIN targets AS st USING(node_type, target_id) + LEFT JOIN buddy_groups AS bg ON st.target_id = bg.s_target_id + AND st.node_type = bg.node_type + LEFT JOIN quota_default_limits AS d USING(id_type, quota_type, pool_id) + LEFT JOIN quota_limits AS l USING(quota_id, id_type, quota_type, pool_id) + WHERE (?1 IS NULL OR e.id_type = ?1) + AND (?2 IS NULL OR e.quota_type = ?2) + AND (?3 IS NULL OR st.pool_id = ?3) + AND (bg.quota_accounting IS NULL OR bg.quota_accounting = 2) + GROUP BY e.quota_id, e.id_type, e.quota_type, st.pool_id + HAVING SUM(e.value) > COALESCE(l.value, d.value)" +); + /// Calculates and pushes exceeded quota info to the nodes pub(crate) async fn distribute_exceeded(app: &impl App) -> Result<()> { if !app.static_info().user_config.quota_enforce { @@ -358,19 +383,8 @@ pub(crate) async fn distribute_exceeded(app: &impl App) -> Result<()> { if quota_licensed { // Fill the prepared messages with matching exceeded quota ids - let mut stmt = tx.prepare_cached(sql!( - "SELECT DISTINCT e.quota_id, e.id_type, e.quota_type, st.pool_id - FROM quota_usage AS e - INNER JOIN targets AS st USING(node_type, target_id) - LEFT JOIN buddy_groups AS bg ON st.target_id = bg.s_target_id - AND st.node_type = bg.node_type - LEFT JOIN quota_default_limits AS d USING(id_type, quota_type, pool_id) - LEFT JOIN quota_limits AS l USING(quota_id, id_type, quota_type, pool_id) - WHERE bg.quota_accounting IS NULL OR bg.quota_accounting = ?1 - GROUP BY e.quota_id, e.id_type, e.quota_type, st.pool_id - HAVING SUM(e.value) > COALESCE(l.value, d.value)" - ))?; - let mut rows = stmt.query([BuddyGroupQuotaAccounting::Both.sql_variant()])?; + let mut stmt = tx.prepare_cached(EXCEEDED_QUOTA_IDS_SQL)?; + let mut rows = stmt.query(params![None::, None::, None::])?; while let Some(row) = rows.next()? { for m in &mut msges { if row.get::<_, PoolId>(3)? == m.pool_id @@ -386,11 +400,11 @@ pub(crate) async fn distribute_exceeded(app: &impl App) -> Result<()> { // If quota is unlicensed, make sure the exceeding ids are removed from the servers. // Otherwise exceeded ids could stay exceeded forever if quota was used before. log::info!( - "Quota enforcement enabled but feature not licensed. Removing quota limits from nodes" + "Quota enforcement enabled but feature not licensed. Removing quota limits \ + from nodes" ); } - // Get all node uids to send the messages to let nodes: Vec = tx.query_map_collect( sql!("SELECT node_uid FROM nodes WHERE node_type IN (?1,?2)"), @@ -467,13 +481,15 @@ fn try_read_quota_ids(path: &Path, read_into: &mut HashSet) -> Result<( mod test { use crate::Config; use crate::app::test::*; - use crate::types::SqliteEnumExt; + use crate::types::{BuddyGroupQuotaAccounting, SqliteEnumExt}; use shared::bee_msg::OpsErr; use shared::bee_msg::quota::{ GetQuotaInfo, GetQuotaInfoResp, QuotaEntry, QuotaInodeSupport, QuotaQueryType, SetExceededQuota, SetExceededQuotaResp, }; use shared::types::{QuotaIdType, QuotaType}; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; #[tokio::test] async fn update() { @@ -647,15 +663,27 @@ mod test { #[tokio::test] async fn distribute_exceeded() { - // This fn doesn't need special config - let app = TestApp::new().await; + // EXCEEDED_QUOTA_IDS_SQL hardcodes this value, it must not silently change + assert_eq!(BuddyGroupQuotaAccounting::Both.sql_variant(), 2); + + // Without both of these, distribute_exceeded() returns early and nothing is asserted + let app = TestApp::with_config(Config { + quota_enable: true, + quota_enforce: true, + ..Default::default() + }) + .await; + + let msg_count = Arc::new(AtomicUsize::new(0)); + let handler_count = msg_count.clone(); app.set_request_handler(move |req| { + handler_count.fetch_add(1, Ordering::SeqCst); let r = req.downcast_ref::().unwrap(); match (r.pool_id, r.id_type, r.quota_type) { (1, QuotaIdType::User, QuotaType::Space) => { - assert_eq!(r.exceeded_quota_ids.as_slice(), &[2, 4, 10]) + assert_eq!(r.exceeded_quota_ids.as_slice(), &[2, 4, 10, 51]) } (1, QuotaIdType::Group, QuotaType::Space) => { assert_eq!(r.exceeded_quota_ids.as_slice(), &[2, 4, 11]) @@ -680,5 +708,11 @@ mod test { }); super::distribute_exceeded(&app).await.unwrap(); + + // Guards against the assertions above silently not running at all + assert!( + msg_count.load(Ordering::SeqCst) > 0, + "no SetExceededQuota messages were sent" + ); } } From 8038a79161cce02c2b894c0680e3721617213fd2 Mon Sep 17 00:00:00 2001 From: Rusty Bee <145002912+rustybee42@users.noreply.github.com> Date: Tue, 2 Jun 2026 12:06:33 +0200 Subject: [PATCH 6/9] fix: Per message setting of response time limit * Add a new associated const to the Msg trait that defines how long to wait for this message when receiving it as a response. To achieve this, make the timeout setting for stream writes and reads a function argument and pass the message specific time limit. * Set/increase the default to 5s as there is no real disadvantage to wait a bit longer for responses, just in case. * For GetQuotaInfoResp, set this to 1m. Querying the maximum of ~167k quota entries per request can take some seconds, according to some quick benchmarks. This should be plenty of room for slower systems/disks. * When reading a response, also use the per-message time limit for the body * Separately set the connect time limit to 2s to avoid waiting too long for fallbacks * Increase the "generic" stream timeouts to 5s. * Limit the concurrent quota requests to half the maximum outgoing connections. This prevents quota requests eating up all the connections and blocking other operations if responses take a long time to arrive. * Make sure the connection limit setting is at least 1 * Renaming to match the term "time_limit" --- mgmtd/src/config.rs | 4 +++ mgmtd/src/quota.rs | 25 ++++++++++++++- shared/src/bee_msg.rs | 3 ++ shared/src/bee_msg/quota.rs | 3 ++ shared/src/conn.rs | 8 +++++ shared/src/conn/incoming.rs | 9 ++++-- shared/src/conn/msg_dispatch.rs | 5 ++- shared/src/conn/outgoing.rs | 56 ++++++++++++++++++++++----------- shared/src/conn/stream.rs | 24 ++++++-------- 9 files changed, 100 insertions(+), 37 deletions(-) diff --git a/mgmtd/src/config.rs b/mgmtd/src/config.rs index a1b5645d..67674722 100644 --- a/mgmtd/src/config.rs +++ b/mgmtd/src/config.rs @@ -451,6 +451,10 @@ impl Config { bail!("Provided file system UUID is not a valid v4 UUID"); } + if self.connection_limit < 1 { + bail!("Connection limit cannot be smaller than 1"); + } + if self.quota_enforce && !self.quota_enable { bail!("Quota enforcement requires quota being enabled"); } diff --git a/mgmtd/src/quota.rs b/mgmtd/src/quota.rs index f159c700..caebb7cd 100644 --- a/mgmtd/src/quota.rs +++ b/mgmtd/src/quota.rs @@ -14,9 +14,11 @@ use shared::bee_msg::quota::{ use shared::types::{NodeType, PoolId, QuotaId, QuotaIdType, QuotaType, TargetId, Uid}; use sqlite::TransactionExt; use sqlite_check::sql; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::path::Path; +use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; +use tokio::sync::Semaphore; use tokio::task::JoinHandle; use tokio::time::Instant; @@ -184,6 +186,10 @@ async fn create_and_send_requests( let mut tasks = vec![]; + // These bound the concurrent requests going on to one node so these potentially long-running + // requests don't block all available connections + let mut semaphores = HashMap::new(); + // Sends one request per (target, id_type, list|range|all) to the respective owner node // Requesting is done concurrently for multiple targets but serialized for the different fetch // modes. @@ -193,10 +199,27 @@ async fn create_and_send_requests( let group_list = group_list.clone(); let user_range = config.quota_user_ids_range.clone(); let group_range = config.quota_group_ids_range.clone(); + let semaphore = semaphores + .entry(t.node_uid) + .or_insert_with(|| Arc::new(Semaphore::new((config.connection_limit / 2).max(1)))) + .clone(); tasks.push(tokio::spawn(async move { let mut responses = vec![]; + let _permit = match semaphore.acquire().await { + Ok(p) => p, + Err(err) => { + log::error!( + "Acquiring permit for fetching quota info for storage target {} from node \ + with uid {} failed: {err:#}", + t.target_id, t.node_uid + ); + + return (t, None); + } + }; + if user_use_all { // Request all entries if no specific ids are configured let resp: Result = app diff --git a/shared/src/bee_msg.rs b/shared/src/bee_msg.rs index d176d0e4..5a634051 100644 --- a/shared/src/bee_msg.rs +++ b/shared/src/bee_msg.rs @@ -6,6 +6,7 @@ use anyhow::{Context, Result, anyhow}; use bee_serde_derive::BeeSerde; use std::any::Any; use std::collections::{HashMap, HashSet}; +use std::time::Duration; pub mod buddy_group; pub mod misc; @@ -26,6 +27,8 @@ pub trait BaseMsg: Any + std::fmt::Debug + Send + Sync + 'static {} pub trait Msg: BaseMsg + Default + Clone { /// Message type as defined in NetMessageTypes.h const ID: MsgId; + /// How long to wait to receive this message as a response + const RESPONSE_TIME_LIMIT: Duration = Duration::from_secs(5); } impl BaseMsg for M where M: Msg {} diff --git a/shared/src/bee_msg/quota.rs b/shared/src/bee_msg/quota.rs index 360ccc88..1413d7f4 100644 --- a/shared/src/bee_msg/quota.rs +++ b/shared/src/bee_msg/quota.rs @@ -135,6 +135,9 @@ pub struct GetQuotaInfoResp { impl Msg for GetQuotaInfoResp { const ID: MsgId = 2098; + // Generously increase the response timeout since querying the maximum of ~167k quota entries + // can take some seconds on slow systems. + const RESPONSE_TIME_LIMIT: Duration = Duration::from_mins(1); } /// Sets exceeded quota information on server nodes. diff --git a/shared/src/conn.rs b/shared/src/conn.rs index 559da65d..75cd5779 100644 --- a/shared/src/conn.rs +++ b/shared/src/conn.rs @@ -1,5 +1,7 @@ //! Connection to other BeeGFS nodes +use std::time::Duration; + mod async_queue; pub mod incoming; pub mod msg_dispatch; @@ -16,3 +18,9 @@ const TCP_BUF_LEN: usize = 4 * 1024 * 1024; /// Must match the `DGRAMMR_(RECV|SEND)BUF_SIZE` value in `DatagramListener.*` in the C/C++ /// codebase. Must be smaller than TCP_BUF_LEN; const UDP_BUF_LEN: usize = 65536; + +/// Reasonable time limit for most stream operations. Notable exceptions are waiting for responses +/// and connecting a stream. +const GENERIC_STREAM_TIME_LIMIT: Duration = Duration::from_secs(5); +/// Short timeout for connecting so the next nic can be tried quickly if this one doesn't work. +const CONNECT_STREAM_TIME_LIMIT: Duration = Duration::from_secs(2); diff --git a/shared/src/conn/incoming.rs b/shared/src/conn/incoming.rs index 45ed69b5..2a60f750 100644 --- a/shared/src/conn/incoming.rs +++ b/shared/src/conn/incoming.rs @@ -139,7 +139,9 @@ async fn read_stream( stream_authentication_required: bool, ) -> Result<()> { // Read header - stream.read_exact(&mut buf[0..Header::LEN]).await?; + stream + .read_exact(&mut buf[0..Header::LEN], GENERIC_STREAM_TIME_LIMIT) + .await?; let header = deserialize_header(&buf[0..Header::LEN])?; @@ -156,7 +158,10 @@ async fn read_stream( // Read body stream - .read_exact(&mut buf[Header::LEN..header.msg_len()]) + .read_exact( + &mut buf[Header::LEN..header.msg_len()], + GENERIC_STREAM_TIME_LIMIT, + ) .await?; // Forward to the dispatcher. The dispatcher is responsible for deserializing, dispatching to diff --git a/shared/src/conn/msg_dispatch.rs b/shared/src/conn/msg_dispatch.rs index 4bf78242..ee64b799 100644 --- a/shared/src/conn/msg_dispatch.rs +++ b/shared/src/conn/msg_dispatch.rs @@ -3,6 +3,7 @@ use super::stream::Stream; use crate::bee_msg::{Header, Msg, deserialize_body, serialize}; use crate::bee_serde::{Deserializable, Serializable}; +use crate::conn::GENERIC_STREAM_TIME_LIMIT; use anyhow::Result; use std::fmt::Debug; use std::future::Future; @@ -40,7 +41,9 @@ pub struct StreamRequest<'a> { impl Request for StreamRequest<'_> { async fn respond(self, msg: &M) -> Result<()> { let msg_len = serialize(msg, self.buf)?; - self.stream.write_all(&self.buf[0..msg_len]).await + self.stream + .write_all(&self.buf[0..msg_len], GENERIC_STREAM_TIME_LIMIT) + .await } fn authenticate_connection(&mut self) { diff --git a/shared/src/conn/outgoing.rs b/shared/src/conn/outgoing.rs index 78fcd2cd..066d2819 100644 --- a/shared/src/conn/outgoing.rs +++ b/shared/src/conn/outgoing.rs @@ -3,9 +3,9 @@ use super::store::Store; use crate::bee_msg::misc::AuthenticateChannel; use crate::bee_msg::{Header, Msg, deserialize_body, deserialize_header, serialize}; use crate::bee_serde::{Deserializable, Serializable}; -use crate::conn::TCP_BUF_LEN; use crate::conn::store::StoredStream; use crate::conn::stream::Stream; +use crate::conn::{CONNECT_STREAM_TIME_LIMIT, GENERIC_STREAM_TIME_LIMIT, TCP_BUF_LEN}; use crate::types::{AuthSecret, Uid}; use anyhow::{Context, Result, bail}; use std::fmt::Debug; @@ -58,7 +58,9 @@ impl Pool { let mut buf = self.store.pop_buf_or_create(); let msg_len = serialize(msg, &mut buf)?; - let resp_header = self.comm_stream(node_uid, &mut buf, msg_len, true).await?; + let resp_header = self + .comm_stream(node_uid, &mut buf, msg_len, Some(R::RESPONSE_TIME_LIMIT)) + .await?; let resp_msg = deserialize_body(&resp_header, &buf[Header::LEN..])?; self.store.push_buf(buf); @@ -75,7 +77,7 @@ impl Pool { let mut buf = self.store.pop_buf_or_create(); let msg_len = serialize(msg, &mut buf)?; - self.comm_stream(node_uid, &mut buf, msg_len, false).await?; + self.comm_stream(node_uid, &mut buf, msg_len, None).await?; self.store.push_buf(buf); @@ -94,19 +96,21 @@ impl Pool { /// 2. Get a permit that allows opening a new stream. Try to open a new stream using the /// available addresses. /// 3. Pop an open stream from the store, waiting until one gets available. + /// + /// If `response_time_limit` is set to `Some(t)`, a response is expected. async fn comm_stream( &self, node_uid: Uid, buf: &mut [u8], send_len: usize, - expect_response: bool, + response_time_limit: Option, ) -> Result
{ debug_assert_eq!(buf.len(), TCP_BUF_LEN); // 1. Pop open streams until communication succeeds or none are left while let Some(stream) = self.store.try_pop_stream(node_uid) { match self - .write_and_read_stream(buf, stream, send_len, expect_response) + .write_and_read_stream(buf, stream, send_len, response_time_limit) .await { Ok(header) => return Ok(header), @@ -132,7 +136,7 @@ impl Pool { continue; } - match Stream::connect_tcp(addr).await { + match Stream::connect_tcp(addr, CONNECT_STREAM_TIME_LIMIT).await { Ok(stream) => { let mut stream = StoredStream::from_stream(stream, permit); @@ -152,7 +156,7 @@ impl Pool { stream .as_mut() - .write_all(&auth_buf[0..msg_len]) + .write_all(&auth_buf[0..msg_len], GENERIC_STREAM_TIME_LIMIT) .await .with_context(err_context)?; @@ -162,7 +166,7 @@ impl Pool { // Communication using the newly opened stream should usually not fail. If // it does, abort. It might be better to just try the next address though. let resp_header = self - .write_and_read_stream(buf, stream, send_len, expect_response) + .write_and_read_stream(buf, stream, send_len, response_time_limit) .await .with_context(err_context)?; @@ -181,15 +185,19 @@ impl Pool { ) } - // 3. Wait for an already open stream becoming available - let stream = timeout(Duration::from_secs(2), self.store.pop_stream(node_uid)) + // 3. Wait for an already open stream becoming available. The timeout is intentionally + // chosen short to avoid big pile up of waiting requests but rather fail quickly. The + // user can always increase the connection limit to work around it. It's also in the + // responsibility of requesters to limit potentially long running requests (e.g. quota + // queries) to not block the whole pool. + let stream = timeout(GENERIC_STREAM_TIME_LIMIT, self.store.pop_stream(node_uid)) .await .map_err(|_| { anyhow::anyhow!("Popping a stream for node with uid {node_uid:?} timed out") })?; let resp_header = self - .write_and_read_stream(buf, stream, send_len, expect_response) + .write_and_read_stream(buf, stream, send_len, response_time_limit) .await .with_context(|| { format!("Communication using existing stream to node with uid {node_uid} failed") @@ -199,26 +207,36 @@ impl Pool { } /// Writes data to the given stream, optionally receives a response and pushes the stream to - /// the store + /// the store. Receives a response if `response_time_limit` is `Some(t)`. async fn write_and_read_stream( &self, buf: &mut [u8], mut stream: StoredStream, send_len: usize, - expect_response: bool, + response_time_limit: Option, ) -> Result
{ - stream.as_mut().write_all(&buf[0..send_len]).await?; + stream + .as_mut() + .write_all(&buf[0..send_len], GENERIC_STREAM_TIME_LIMIT) + .await?; - let header = if expect_response { - // Read header - stream.as_mut().read_exact(&mut buf[0..Header::LEN]).await?; + let header = if let Some(tl) = response_time_limit { + // Read header - wait for the per-message defined time limit. + stream + .as_mut() + .read_exact(&mut buf[0..Header::LEN], tl) + .await?; let header = deserialize_header(&buf[0..Header::LEN])?; - // Read body + // Read body - the header has already been received, so the body should follow + // immediately as currently nodes serialize whole messages before sending. Still + // choosing the (potentially higher) per-message limit in case that changes at some + // point. stream .as_mut() - .read_exact(&mut buf[Header::LEN..header.msg_len()]) + .read_exact(&mut buf[Header::LEN..header.msg_len()], tl) .await?; + header } else { Header::default() diff --git a/shared/src/conn/stream.rs b/shared/src/conn/stream.rs index a98f88fe..b8b0713f 100644 --- a/shared/src/conn/stream.rs +++ b/shared/src/conn/stream.rs @@ -9,8 +9,6 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpStream; use tokio::time::timeout; -const TIMEOUT: Duration = Duration::from_secs(2); - /// A connected generic stream. /// /// Provides functionality to communicate with the connected peer. Can support multiple @@ -39,9 +37,9 @@ impl From for Stream { impl Stream { /// Connect to peer using TCP and obtain a [Stream] object. /// - /// Times out after [TIMEOUT]. - pub async fn connect_tcp(addr: &SocketAddr) -> Result { - let stream = match timeout(TIMEOUT, TcpStream::connect(addr)).await { + /// Times out after `time_limit`. + pub async fn connect_tcp(addr: &SocketAddr, time_limit: Duration) -> Result { + let stream = match timeout(time_limit, TcpStream::connect(addr)).await { Ok(res) => res?, Err(_) => bail!("Connecting a TCP stream to {addr} timed out"), }; @@ -74,13 +72,11 @@ impl Stream { /// Reads from the stream into the provided buffer. /// /// The buffer will be filled completely before the future completes. Times out after - /// [TIMEOUT]. + /// `time_limit`. /// /// **Important**: Not cancel safe. If a timeout occurs, the stream may not be reused. - // Clippy: Suppress false positive - #[allow(clippy::needless_pass_by_ref_mut)] - pub async fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> { - match timeout(TIMEOUT, async { + pub async fn read_exact(&mut self, buf: &mut [u8], time_limit: Duration) -> Result<()> { + match timeout(time_limit, async { match &mut self.stream { InnerStream::Tcp(s) => { s.read_exact(buf).await?; @@ -98,11 +94,11 @@ impl Stream { /// Writes to the stream from the provided buffer. /// /// The buffer will be written completely before the future completes. Times out after - /// [TIMEOUT]. + /// `time_limit`. /// /// **Important**: Not cancel safe. If a timeout occurs, the stream may not be reused. - pub async fn write_all(&mut self, buf: &[u8]) -> Result<()> { - match timeout(TIMEOUT, async { + pub async fn write_all(&mut self, buf: &[u8], time_limit: Duration) -> Result<()> { + match timeout(time_limit, async { match &mut self.stream { InnerStream::Tcp(s) => { s.write_all(buf).await?; @@ -113,7 +109,7 @@ impl Stream { .await { Ok(res) => res, - Err(_) => Err(anyhow!("Writing to a stream to {} timed out", self.addr())), + Err(_) => Err(anyhow!("Writing to stream to {} timed out", self.addr())), } } From d4d67b2c8d46b54390c40624cef42b129b7b9648 Mon Sep 17 00:00:00 2001 From: Rusty Bee <145002912+rustybee42@users.noreply.github.com> Date: Mon, 17 Aug 2026 13:27:40 +0200 Subject: [PATCH 7/9] fix: various issues in quota.rs * fetch_and_update() doesn't error out anymore on inner errors, always processing all tasks. This prevents running tasks getting detached and potentially keep running during the next fetch, allowing more concurrent connections being used for quota than the semaphore allows. * Error out if any ids were configured (no all mode) but none were resolved (e.g. by providing an empty id file) * Fix various small issues and redundancies * Add some commentary --- mgmtd/src/quota.rs | 143 +++++++++++++++++++++++++-------------------- 1 file changed, 81 insertions(+), 62 deletions(-) diff --git a/mgmtd/src/quota.rs b/mgmtd/src/quota.rs index caebb7cd..b533e7ed 100644 --- a/mgmtd/src/quota.rs +++ b/mgmtd/src/quota.rs @@ -5,7 +5,7 @@ mod system_id; use crate::app::*; use crate::license::LicensedFeature; use crate::types::SqliteEnumExt; -use anyhow::{Context as AnyhowContext, Result}; +use anyhow::{Context as AnyhowContext, Result, bail}; use rusqlite::params; use shared::bee_msg::OpsErr; use shared::bee_msg::quota::{ @@ -17,7 +17,6 @@ use sqlite_check::sql; use std::collections::{HashMap, HashSet}; use std::path::Path; use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; use tokio::sync::Semaphore; use tokio::task::JoinHandle; use tokio::time::Instant; @@ -43,8 +42,7 @@ pub(crate) async fn fetch_and_update(app: &impl App) -> Result<()> { sql!( "SELECT target_id, pool_id, node_uid FROM storage_targets - INNER JOIN nodes USING(node_type, node_id) - WHERE node_id IS NOT NULL" + INNER JOIN nodes USING(node_type, node_id)" ), [], |row| { @@ -64,66 +62,80 @@ pub(crate) async fn fetch_and_update(app: &impl App) -> Result<()> { } let targets_to_query_count = targets_to_query.len(); - let start_time = Instant::now(); - let entry_counter = AtomicUsize::new(0); let tasks = create_and_send_requests(app, targets_to_query).await?; // Await all the responses - for t in tasks { - let (target, entries) = t.await?; - - // Only process that target if there were not errors when fetching for this target - if let Some(entries) = entries { - entry_counter.fetch_add(entries.len(), Ordering::Relaxed); - - app.write_tx(move |tx| { - // Always delete all the old entries for that target to make sure entries for no - // longer queried ids are removed. We always get the complete list from the - // storages and we only update if there was no fetch error. - tx.execute_cached( - sql!("DELETE FROM quota_usage WHERE target_id = ?1"), - [target.target_id], - )?; + let mut entry_counter = 0; + for (target, jh) in tasks { + let res = async { + let entries = jh.await?; + + // Only process that target if there were not errors when fetching for this target + if let Some(entries) = entries { + entry_counter += entries.len(); + + app.write_tx(move |tx| { + // Always delete all the old entries for that target to make sure entries for no + // longer queried ids are removed. We always get the complete list from the + // storages and we only update if there was no fetch error. + // There is one task per target with merged results from multiple queries, so no + // accidental override here. + tx.execute_cached( + sql!("DELETE FROM quota_usage WHERE target_id = ?1"), + [target.target_id], + )?; + + // The entry list can contain duplicated entries if both range and list mode + // are configured as they use two separate requests, thus the OR IGNORE. + let mut insert_stmt = tx.prepare_cached(sql!( + "INSERT OR IGNORE + INTO quota_usage (quota_id, id_type, quota_type, target_id, value) + VALUES (?1, ?2, ?3 ,?4 ,?5)" + ))?; + + for e in entries { + if e.space > 0 { + insert_stmt.execute(params![ + e.id, + e.id_type.sql_variant(), + QuotaType::Space.sql_variant(), + target.target_id, + e.space + ])?; + } - let mut insert_stmt = tx.prepare_cached(sql!( - "INSERT OR IGNORE - INTO quota_usage (quota_id, id_type, quota_type, target_id, value) - VALUES (?1, ?2, ?3 ,?4 ,?5)" - ))?; - - for e in entries { - if e.space > 0 { - insert_stmt.execute(params![ - e.id, - e.id_type.sql_variant(), - QuotaType::Space.sql_variant(), - target.target_id, - e.space - ])?; + if e.inodes > 0 { + insert_stmt.execute(params![ + e.id, + e.id_type.sql_variant(), + QuotaType::Inode.sql_variant(), + target.target_id, + e.inodes + ])?; + } } - if e.inodes > 0 { - insert_stmt.execute(params![ - e.id, - e.id_type.sql_variant(), - QuotaType::Inode.sql_variant(), - target.target_id, - e.inodes - ])?; - } - } + Ok(()) + }) + .await?; + } - Ok(()) - }) - .await?; + Ok(()) as Result<_> + } + .await; + + if let Err(err) = res { + log::error!( + "Receiving and storing quota info from storage target {} failed: {err:#}", + target.target_id + ); } } log::info!( - "Fetched and stored {} quota entries from {} targets in {:?}", - entry_counter.load(Ordering::Relaxed), + "Fetched and stored {entry_counter} quota entries from {} targets in {:?}", targets_to_query_count, start_time.elapsed() ); @@ -136,7 +148,7 @@ pub(crate) async fn fetch_and_update(app: &impl App) -> Result<()> { async fn create_and_send_requests( app: &impl App, targets: Vec, -) -> Result>)>>> { +) -> Result>>)>> { let config = &app.static_info().user_config; // The to-be-queried IDs @@ -184,7 +196,14 @@ async fn create_and_send_requests( && config.quota_group_ids_file.is_none() && config.quota_group_ids_range.is_none(); - let mut tasks = vec![]; + if !user_use_all && user_list.is_empty() && config.quota_user_ids_range.is_none() { + bail!("User quota ID selection is configured but resolved to no IDs"); + } + if !group_use_all && group_list.is_empty() && config.quota_group_ids_range.is_none() { + bail!("Group quota ID selection is configured but resolved to no IDs"); + } + + let mut tasks: Vec<(TargetToQuery, JoinHandle>)> = vec![]; // These bound the concurrent requests going on to one node so these potentially long-running // requests don't block all available connections @@ -204,7 +223,7 @@ async fn create_and_send_requests( .or_insert_with(|| Arc::new(Semaphore::new((config.connection_limit / 2).max(1)))) .clone(); - tasks.push(tokio::spawn(async move { + tasks.push((t, tokio::spawn(async move { let mut responses = vec![]; let _permit = match semaphore.acquire().await { @@ -216,7 +235,7 @@ async fn create_and_send_requests( t.target_id, t.node_uid ); - return (t, None); + return None; } }; @@ -308,9 +327,8 @@ async fn create_and_send_requests( } } - let results = extract_results(&t, responses); - (t, results) - })); + extract_results(&t, responses) + }))); } Ok(tasks) @@ -340,7 +358,7 @@ fn extract_results( } else { log::error!( "Fetching quota info for storage target {} from node with uid \ -{} failed:{errs}", + {} failed:{errs}", target.target_id, target.node_uid ); @@ -359,7 +377,7 @@ fn extract_results( /// neither the fixed nor the optional form of the id type / quota type filters can seek on it (thus /// no difference in performance). pub(crate) const EXCEEDED_QUOTA_IDS_SQL: &str = sql!( - "SELECT DISTINCT e.quota_id, e.id_type, e.quota_type, st.pool_id + "SELECT e.quota_id, e.id_type, e.quota_type, st.pool_id FROM quota_usage AS e INNER JOIN targets AS st USING(node_type, target_id) LEFT JOIN buddy_groups AS bg ON st.target_id = bg.s_target_id @@ -492,9 +510,10 @@ pub(crate) async fn distribute_exceeded(app: &impl App) -> Result<()> { /// /// IDs must be in numerical form and separated by any whitespace. fn try_read_quota_ids(path: &Path, read_into: &mut HashSet) -> Result<()> { - let data = std::fs::read_to_string(path)?; + let data = std::fs::read_to_string(path) + .with_context(|| format!("Could not read quota id file {path:?}"))?; for id in data.split_whitespace().map(|e| e.parse()) { - read_into.insert(id.context("Invalid syntax in quota file {path}")?); + read_into.insert(id.with_context(|| format!("Invalid syntax in quota id file {path:?}"))?); } Ok(()) From 4737a4b3e37000c31c9769d218e29090dc11e2f9 Mon Sep 17 00:00:00 2001 From: Rusty Bee <145002912+rustybee42@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:58:19 +0200 Subject: [PATCH 8/9] fix: on receive, check message fits into the provided buffer before receiving --- shared/src/bee_msg.rs | 15 ++++++++++++++- shared/src/conn/incoming.rs | 4 ++-- shared/src/conn/outgoing.rs | 2 +- 3 files changed, 17 insertions(+), 4 deletions(-) diff --git a/shared/src/bee_msg.rs b/shared/src/bee_msg.rs index 5a634051..1716b7c8 100644 --- a/shared/src/bee_msg.rs +++ b/shared/src/bee_msg.rs @@ -155,6 +155,9 @@ pub fn serialize(msg: &M, buf: &mut [u8]) -> Result Result
{ @@ -184,6 +187,16 @@ pub fn deserialize_header(buf: &[u8]) -> Result
{ .context(CTX); } + if header.msg_len as usize > buf.len() { + return Err(anyhow!( + "Received BeeMsg doesn't fit into the provided buffer: Reported length {}, \ + buffer size is {}", + header.msg_len, + buf.len() + )) + .context(CTX); + } + Ok(header) } @@ -209,7 +222,7 @@ pub fn deserialize_body(header: &Header, buf: &[u8]) -> /// # Return value /// Returns the deserialized message. pub fn deserialize(buf: &[u8]) -> Result { - let header = deserialize_header(&buf[0..Header::LEN])?; + let header = deserialize_header(buf)?; let msg = deserialize_body(&header, &buf[Header::LEN..])?; Ok(msg) } diff --git a/shared/src/conn/incoming.rs b/shared/src/conn/incoming.rs index 2a60f750..f1bbf29b 100644 --- a/shared/src/conn/incoming.rs +++ b/shared/src/conn/incoming.rs @@ -143,7 +143,7 @@ async fn read_stream( .read_exact(&mut buf[0..Header::LEN], GENERIC_STREAM_TIME_LIMIT) .await?; - let header = deserialize_header(&buf[0..Header::LEN])?; + let header = deserialize_header(buf)?; // check authentication if stream_authentication_required @@ -235,7 +235,7 @@ async fn recv_datagram(sock: Arc, msg_handler: impl DispatchRequest) // immediately tokio::spawn(async move { if let Err(err) = async { - let header = deserialize_header(&buf[0..Header::LEN])?; + let header = deserialize_header(&buf)?; let req = SocketRequest { sock, diff --git a/shared/src/conn/outgoing.rs b/shared/src/conn/outgoing.rs index 066d2819..e824915a 100644 --- a/shared/src/conn/outgoing.rs +++ b/shared/src/conn/outgoing.rs @@ -226,7 +226,7 @@ impl Pool { .as_mut() .read_exact(&mut buf[0..Header::LEN], tl) .await?; - let header = deserialize_header(&buf[0..Header::LEN])?; + let header = deserialize_header(buf)?; // Read body - the header has already been received, so the body should follow // immediately as currently nodes serialize whole messages before sending. Still From 0aa879cd50050348409ab166ded63a81d1a4a2e2 Mon Sep 17 00:00:00 2001 From: Rusty Bee <145002912+rustybee42@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:25:21 +0200 Subject: [PATCH 9/9] feat/fix: Query more quota entries if storage signals there are more This works using a compat feature flag in the response message. Run the requests in a loop until there are no more entries left. Use the range query for all mode to be able to chunk the requests. To be able to read the flag, add a request_with_header() method. --- mgmtd/src/app.rs | 9 +- mgmtd/src/app/runtime.rs | 10 ++- mgmtd/src/app/test.rs | 8 ++ mgmtd/src/quota.rs | 167 +++++++++++++++++++++++------------- shared/src/bee_msg/quota.rs | 17 +--- shared/src/conn/outgoing.rs | 6 +- 6 files changed, 138 insertions(+), 79 deletions(-) diff --git a/mgmtd/src/app.rs b/mgmtd/src/app.rs index 203809f5..6ce89d9f 100644 --- a/mgmtd/src/app.rs +++ b/mgmtd/src/app.rs @@ -10,7 +10,7 @@ use anyhow::Result; use protobuf::license::GetCertDataResult; pub(crate) use runtime::RuntimeApp; use rusqlite::{Connection, Transaction}; -use shared::bee_msg::Msg; +use shared::bee_msg::{Header, Msg}; use shared::bee_serde::{Deserializable, Serializable}; use shared::types::{NodeId, NodeType, Uid}; use std::fmt::Debug; @@ -50,6 +50,13 @@ pub(crate) trait App: Debug + Clone + Send + 'static { ) -> impl Future> + Send; // BeeMsg communication + // + /// Send a [Msg] to a node via TCP and receive the response + fn request_with_header( + &self, + node_uid: Uid, + msg: &M, + ) -> impl Future> + Send; /// Send a [Msg] to a node via TCP and receive the response fn request( diff --git a/mgmtd/src/app/runtime.rs b/mgmtd/src/app/runtime.rs index a465be1a..6fcf0cdc 100644 --- a/mgmtd/src/app/runtime.rs +++ b/mgmtd/src/app/runtime.rs @@ -116,12 +116,20 @@ impl App for RuntimeApp { Connections::conn(&self.db, op).await } + async fn request_with_header( + &self, + node_uid: Uid, + msg: &M, + ) -> Result<(R, Header)> { + Pool::request(&self.conn, node_uid, msg).await + } + async fn request( &self, node_uid: Uid, msg: &M, ) -> Result { - Pool::request(&self.conn, node_uid, msg).await + Pool::request(&self.conn, node_uid, msg).await.map(|e| e.0) } async fn send_notifications( diff --git a/mgmtd/src/app/test.rs b/mgmtd/src/app/test.rs index f4b61dc6..bd8c7215 100644 --- a/mgmtd/src/app/test.rs +++ b/mgmtd/src/app/test.rs @@ -127,6 +127,14 @@ impl App for TestApp { Connections::conn(&self.db, op).await } + async fn request_with_header( + &self, + node_uid: Uid, + msg: &M, + ) -> Result<(R, Header)> { + Ok((self.request(node_uid, msg).await?, Header::default())) + } + async fn request( &self, _node_uid: Uid, diff --git a/mgmtd/src/quota.rs b/mgmtd/src/quota.rs index b533e7ed..24c342e6 100644 --- a/mgmtd/src/quota.rs +++ b/mgmtd/src/quota.rs @@ -15,6 +15,7 @@ use shared::types::{NodeType, PoolId, QuotaId, QuotaIdType, QuotaType, TargetId, use sqlite::TransactionExt; use sqlite_check::sql; use std::collections::{HashMap, HashSet}; +use std::ops::RangeInclusive; use std::path::Path; use std::sync::Arc; use tokio::sync::Semaphore; @@ -211,19 +212,19 @@ async fn create_and_send_requests( // Sends one request per (target, id_type, list|range|all) to the respective owner node // Requesting is done concurrently for multiple targets but serialized for the different fetch - // modes. - for t in targets { + // modes and multiple chunks. + for target in targets { let app = app.clone(); let user_list = user_list.clone(); let group_list = group_list.clone(); let user_range = config.quota_user_ids_range.clone(); let group_range = config.quota_group_ids_range.clone(); let semaphore = semaphores - .entry(t.node_uid) + .entry(target.node_uid) .or_insert_with(|| Arc::new(Semaphore::new((config.connection_limit / 2).max(1)))) .clone(); - tasks.push((t, tokio::spawn(async move { + tasks.push((target, tokio::spawn(async move { let mut responses = vec![]; let _permit = match semaphore.acquire().await { @@ -231,8 +232,9 @@ async fn create_and_send_requests( Err(err) => { log::error!( "Acquiring permit for fetching quota info for storage target {} from node \ - with uid {} failed: {err:#}", - t.target_id, t.node_uid + with uid {} failed: {err:#}", + target.target_id, + target.node_uid ); return None; @@ -240,25 +242,39 @@ async fn create_and_send_requests( }; if user_use_all { - // Request all entries if no specific ids are configured - let resp: Result = app - .request( - t.node_uid, - &GetQuotaInfo::with_all(QuotaIdType::User, t.target_id, t.pool_id), - ) - .await; - - responses.push(("User all", resp)); + // If configured, query the whole id space + range_requests( + app.clone(), + target.node_uid, + QuotaIdType::User, + &target, + &(0..=QuotaId::MAX), + "User id all", + &mut responses, + ) + .await; } else { // Otherwise query the configured ids via list and range + if let Some(ref range) = user_range { + range_requests( + app.clone(), + target.node_uid, + QuotaIdType::User, + &target, + range, + "User id range", + &mut responses, + ) + .await; + } if !user_list.is_empty() { let resp: Result = app .request( - t.node_uid, + target.node_uid, &GetQuotaInfo::with_list( QuotaIdType::User, - t.target_id, - t.pool_id, + target.target_id, + target.pool_id, user_list, ), ) @@ -266,43 +282,41 @@ async fn create_and_send_requests( responses.push(("User id list", resp)); } - if let Some(ref range) = user_range { - let resp: Result = app - .request( - t.node_uid, - &GetQuotaInfo::with_range( - QuotaIdType::User, - t.target_id, - t.pool_id, - range, - ), - ) - .await; - - responses.push(("User id range", resp)); - } } if group_use_all { - // Request all entries if no specific ids are configured - let resp: Result = app - .request( - t.node_uid, - &GetQuotaInfo::with_all(QuotaIdType::Group, t.target_id, t.pool_id), - ) - .await; - - responses.push(("Group all", resp)); + range_requests( + app.clone(), + target.node_uid, + QuotaIdType::Group, + &target, + &(0..=QuotaId::MAX), + "Group id all", + &mut responses, + ) + .await; } else { // Otherwise query the configured ids via list and range + if let Some(ref range) = group_range { + range_requests( + app.clone(), + target.node_uid, + QuotaIdType::Group, + &target, + range, + "Group id range", + &mut responses, + ) + .await; + } if !group_list.is_empty() { let resp: Result = app .request( - t.node_uid, + target.node_uid, &GetQuotaInfo::with_list( QuotaIdType::Group, - t.target_id, - t.pool_id, + target.target_id, + target.pool_id, group_list, ), ) @@ -310,30 +324,61 @@ async fn create_and_send_requests( responses.push(("Group id list", resp)); } - if let Some(ref range) = group_range { - let resp: Result = app - .request( - t.node_uid, - &GetQuotaInfo::with_range( - QuotaIdType::Group, - t.target_id, - t.pool_id, - range, - ), - ) - .await; - - responses.push(("Group id range", resp)); - } } - extract_results(&t, responses) + extract_results(&target, responses) }))); } Ok(tasks) } +async fn range_requests( + app: impl App, + node_uid: Uid, + id_type: QuotaIdType, + target: &TargetToQuery, + range: &RangeInclusive, + log_str: &'static str, + responses: &mut Vec<(&str, Result)>, +) { + let mut range_start = *range.start(); + let range_end = *range.end(); + let mut has_more = true; + + while has_more && range_start <= range_end { + let resp = app + .request_with_header::<_, GetQuotaInfoResp>( + node_uid, + &GetQuotaInfo::with_range( + id_type, + target.target_id, + target.pool_id, + &(range_start..=range_end), + ), + ) + .await; + + (has_more, range_start) = resp + .as_ref() + .map(|e| { + let range_start = + e.0.quota_entry + .last() + .map(|s| s.id.saturating_add(1)) + .unwrap_or_default(); + let has_more = range_start > 0 + && e.1.msg_compat_feature_flags & GetQuotaInfoResp::HAS_MORE_ENTRIES_COMPATFLAG + != 0; + + (has_more, range_start) + }) + .unwrap_or_default(); + + responses.push((log_str, resp.map(|e| e.0))); + } +} + /// Extracts the quota entries from the response message or log the errors fn extract_results( target: &TargetToQuery, diff --git a/shared/src/bee_msg/quota.rs b/shared/src/bee_msg/quota.rs index 1413d7f4..f02a8639 100644 --- a/shared/src/bee_msg/quota.rs +++ b/shared/src/bee_msg/quota.rs @@ -58,19 +58,6 @@ impl GetQuotaInfo { pool_id, } } - - pub fn with_all(id_type: QuotaIdType, target_id: TargetId, pool_id: PoolId) -> Self { - Self { - query_type: QuotaQueryType::All, - id_type, - id_range_start: 0, - id_range_end: 0, - id_list: vec![], - transfer_method: GetQuotaInfoTransferMethod::AllTargetsOneRequestPerTarget, - target_id, - pool_id, - } - } } impl Msg for GetQuotaInfo { @@ -133,6 +120,10 @@ pub struct GetQuotaInfoResp { pub quota_entry: Vec, } +impl GetQuotaInfoResp { + pub const HAS_MORE_ENTRIES_COMPATFLAG: u8 = 1; +} + impl Msg for GetQuotaInfoResp { const ID: MsgId = 2098; // Generously increase the response timeout since querying the maximum of ~167k quota entries diff --git a/shared/src/conn/outgoing.rs b/shared/src/conn/outgoing.rs index e824915a..6482608e 100644 --- a/shared/src/conn/outgoing.rs +++ b/shared/src/conn/outgoing.rs @@ -47,12 +47,12 @@ impl Pool { } } - /// Sends a [Msg] to a node and receives the response. + /// Send a [Msg] to a node, receive the response and return it together with its header. pub async fn request( &self, node_uid: Uid, msg: &M, - ) -> Result { + ) -> Result<(R, Header)> { log::trace!("REQUEST to {node_uid:?}: {msg:?}"); let mut buf = self.store.pop_buf_or_create(); @@ -67,7 +67,7 @@ impl Pool { log::trace!("RESPONSE RECEIVED from {node_uid:?}: {resp_msg:?}"); - Ok(resp_msg) + Ok((resp_msg, resp_header)) } /// Sends a [Msg] to a node and does **not** receive a response.