Refactor out legacy parts of ingest controller

quickwit/quickwit-control-plane/src/control_plane.rs

@@ -17,6 +17,8 @@ use std::collections::{BTreeMap, BTreeSet};
 use std::fmt;
 use std::fmt::Formatter;
 use std::num::NonZeroUsize;
+use std::sync::Arc;
+use std::sync::atomic::Ordering;
 use std::time::Duration;
 
 use anyhow::Context;
@@ -59,8 +61,8 @@ use crate::IndexerPool;
 use crate::cooldown_map::{CooldownMap, CooldownStatus};
 use crate::debouncer::Debouncer;
 use crate::indexing_scheduler::{IndexingScheduler, IndexingSchedulerState};
-use crate::ingest::IngestController;
-use crate::ingest::ingest_controller::{IngestControllerStats, RebalanceShardsCallback};
+use crate::ingest::ingest_controller::RebalanceShardsCallback;
+use crate::ingest::{IngestController, LegacyScalingController};
 use crate::metrics::{METASTORE_ERROR_ABORTED, METASTORE_ERROR_MAYBE_EXECUTED, RESTART_TOTAL};
 use crate::model::ControlPlaneModel;
 
@@ -97,7 +99,8 @@ pub struct ControlPlane {
     // - the ingest controller is in charge of managing ingesters: it opens and closes shards on
     // the different ingesters.
     indexing_scheduler: IndexingScheduler,
-    ingest_controller: IngestController,
+    ingest_controller: Arc<IngestController>,
+    legacy_scaling_controller: LegacyScalingController,
     metastore: MetastoreServiceClient,
     model: ControlPlaneModel,
     prune_shard_cooldown: CooldownMap<(IndexId, SourceId)>,
@@ -134,9 +137,12 @@ impl ControlPlane {
                     / shared_consts::MIB as f32;
                 let indexing_scheduler =
                     IndexingScheduler::new(cluster_id, self_node_id.clone(), indexer_pool.clone());
-                let ingest_controller = IngestController::new(
+                let ingest_controller = Arc::new(IngestController::new(
                     metastore.clone(),
                     ingester_pool.clone(),
+                ));
+                let legacy_scaling_controller = LegacyScalingController::new(
+                    ingest_controller.clone(),
                     shard_throughput_limit_mib,
                     cluster_config.shard_scale_up_factor,
                 );
@@ -148,6 +154,7 @@ impl ControlPlane {
                     cluster_config: cluster_config.clone(),
                     indexing_scheduler,
                     ingest_controller,
+                    legacy_scaling_controller,
                     metastore: metastore.clone(),
                     model: Default::default(),
                     prune_shard_cooldown: CooldownMap::new(NonZeroUsize::new(1024).unwrap()),
@@ -162,7 +169,7 @@ impl ControlPlane {
 #[derive(Debug, Clone, Serialize, Default)]
 pub struct ControlPlaneObservableState {
     pub indexing_scheduler: IndexingSchedulerState,
-    pub ingest_controller: IngestControllerStats,
+    pub num_rebalance_shards_ops: usize,
     pub num_indexes: usize,
     pub num_sources: usize,
     pub readiness: bool,
@@ -179,7 +186,10 @@ impl Actor for ControlPlane {
     fn observable_state(&self) -> Self::ObservableState {
         ControlPlaneObservableState {
             indexing_scheduler: self.indexing_scheduler.observable_state(),
-            ingest_controller: self.ingest_controller.stats,
+            num_rebalance_shards_ops: self
+                .ingest_controller
+                .num_rebalance_shards_ops
+                .load(Ordering::Relaxed),
             num_indexes: self.model.num_indexes(),
             num_sources: self.model.num_sources(),
             readiness: *self.readiness_tx.borrow(),
@@ -964,7 +974,7 @@ impl Handler<LocalShardsUpdate> for ControlPlane {
         ctx: &ActorContext<Self>,
     ) -> Result<Self::Reply, ActorExitStatus> {
         if let Err(metastore_error) = self
-            .ingest_controller
+            .legacy_scaling_controller
             .handle_local_shards_update(local_shards_update, &mut self.model, ctx.progress())
             .await
         {
@@ -1829,7 +1839,7 @@ mod tests {
         tokio::time::timeout(Duration::from_secs(2), async {
             loop {
                 let state = control_plane_mailbox.ask(Observe).await.unwrap();
-                if state.ingest_controller.num_rebalance_shards_ops >= 2 {
+                if state.num_rebalance_shards_ops >= 2 {
                     break;
                 }
                 tokio::time::sleep(Duration::from_millis(10)).await;
@@ -2597,14 +2607,13 @@ mod tests {
 
         universe.sleep(Duration::from_secs(10)).await;
 
-        let ingest_controller_stats = control_plane_handle
+        let num_rebalance_1 = control_plane_handle
             .process_pending_and_observe()
             .await
             .state_opt
             .as_ref()
             .unwrap()
-            .ingest_controller;
-        let num_rebalance_1 = ingest_controller_stats.num_rebalance_shards_ops;
+            .num_rebalance_shards_ops;
         assert!(num_rebalance_1 >= 1);
 
         ingester_pool_change_tx
@@ -2613,15 +2622,13 @@ mod tests {
 
         universe.sleep(Duration::from_secs(10)).await;
 
-        let ingest_controller_stats = control_plane_handle
+        let num_rebalance_2 = control_plane_handle
             .process_pending_and_observe()
             .await
             .state_opt
             .as_ref()
             .unwrap()
-            .ingest_controller;
-
-        let num_rebalance_2 = ingest_controller_stats.num_rebalance_shards_ops;
+            .num_rebalance_shards_ops;
         assert!(num_rebalance_2 > num_rebalance_1);
 
         universe.assert_quit().await;

quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs

@@ -206,7 +206,7 @@ fn compute_load_per_shard(shard_entries: &[&ShardEntry]) -> NonZeroU32 {
         let num_shards = shard_entries.len().max(1) as u64;
         let average_throughput_per_shard_bytes: u64 = shard_entries
             .iter()
-            .map(|shard_entry| shard_entry.long_term_ingestion_rate.0 as u64 * bytesize::MIB)
+            .map(|shard_entry| shard_entry.long_term_ingestion_rate.as_u64())
             .sum::<u64>()
             .div_ceil(num_shards)
             // A shard throughput cannot exceed PIPELINE_THROUGHPUT in the long term (this is

quickwit/quickwit-control-plane/src/ingest/legacy_scaling_arbiter.rs

@@ -14,9 +14,17 @@
 
 use std::num::NonZeroUsize;
 
-use crate::model::{ScalingMode, ShardStats};
+use crate::model::ShardStats;
 
-pub(crate) struct ScalingArbiter {
+#[derive(Debug, Clone, Copy, Eq, PartialEq)]
+pub(crate) enum ScalingMode {
+    /// Scale up by adding this number of shards
+    Up(usize),
+    /// Scale down by removing one shard
+    Down,
+}
+
+pub(crate) struct LegacyScalingArbiter {
     // Threshold in MiB/s below which we decrease the number of shards.
     scale_down_shards_threshold_mib_per_sec: f32,
 
@@ -36,12 +44,12 @@ pub(crate) struct ScalingArbiter {
     shard_scale_up_factor: f32,
 }
 
-impl ScalingArbiter {
+impl LegacyScalingArbiter {
     pub fn with_max_shard_ingestion_throughput_mib_per_sec(
         max_shard_throughput_mib_per_sec: f32,
         shard_scale_up_factor: f32,
-    ) -> ScalingArbiter {
-        ScalingArbiter {
+    ) -> LegacyScalingArbiter {
+        LegacyScalingArbiter {
             scale_up_shards_short_term_threshold_mib_per_sec: max_shard_throughput_mib_per_sec
                 * 0.8f32,
             scale_up_shards_long_term_threshold_mib_per_sec: max_shard_throughput_mib_per_sec
@@ -116,15 +124,15 @@ impl ScalingArbiter {
 mod tests {
     use std::num::NonZeroUsize;
 
-    use super::ScalingArbiter;
-    use crate::model::{ScalingMode, ShardStats};
+    use super::{LegacyScalingArbiter, ScalingMode};
+    use crate::model::ShardStats;
 
     #[test]
     fn test_scaling_arbiter_one_by_one() {
         // use shard throughput 10MiB to simplify calculations
         // with a factor close to 1 shards are effectively added 1 by 1
         let scaling_arbiter =
-            ScalingArbiter::with_max_shard_ingestion_throughput_mib_per_sec(10.0, 1.01);
+            LegacyScalingArbiter::with_max_shard_ingestion_throughput_mib_per_sec(10.0, 1.01);
         assert_eq!(
             scaling_arbiter.should_scale(
                 ShardStats {
@@ -215,7 +223,7 @@ mod tests {
     fn test_scaling_arbiter_2x() {
         // use shard throughput 10MiB to simplify calculations
         let scaling_arbiter =
-            ScalingArbiter::with_max_shard_ingestion_throughput_mib_per_sec(10.0, 2.);
+            LegacyScalingArbiter::with_max_shard_ingestion_throughput_mib_per_sec(10.0, 2.);
         assert_eq!(
             scaling_arbiter.should_scale(
                 ShardStats {
@@ -319,7 +327,7 @@ mod tests {
     fn test_scale_up_computations() {
         // use shard throughput 10MiB to simplify calculations
         let scaling_arbiter =
-            ScalingArbiter::with_max_shard_ingestion_throughput_mib_per_sec(10.0, 1.5);
+            LegacyScalingArbiter::with_max_shard_ingestion_throughput_mib_per_sec(10.0, 1.5);
 
         let shard_stats = ShardStats {
             num_open_shards: 0,
@@ -400,7 +408,7 @@ mod tests {
     #[test]
     fn test_scaling_arbiter_idle() {
         let scaling_arbiter =
-            ScalingArbiter::with_max_shard_ingestion_throughput_mib_per_sec(10.0, 1.5);
+            LegacyScalingArbiter::with_max_shard_ingestion_throughput_mib_per_sec(10.0, 1.5);
 
         let shard_stats = ShardStats {
             num_open_shards: 0,
@@ -426,7 +434,7 @@ mod tests {
     #[test]
     fn test_scaling_arbiter_min_shards() {
         let scaling_arbiter =
-            ScalingArbiter::with_max_shard_ingestion_throughput_mib_per_sec(10.0, 1.5);
+            LegacyScalingArbiter::with_max_shard_ingestion_throughput_mib_per_sec(10.0, 1.5);
 
         let shard_stats = ShardStats {
             num_open_shards: 1,

quickwit/quickwit-control-plane/src/ingest/mod.rs

@@ -13,8 +13,11 @@
 // limitations under the License.
 
 pub(crate) mod ingest_controller;
-mod scaling_arbiter;
+mod legacy_scaling_arbiter;
+mod legacy_scaling_controller;
 mod wait_handle;
 
 pub use ingest_controller::IngestController;
+pub(crate) use legacy_scaling_arbiter::ScalingMode;
+pub(crate) use legacy_scaling_controller::LegacyScalingController;
 pub use wait_handle::WaitHandle;

quickwit/quickwit-control-plane/src/model/mod.rs

@@ -36,9 +36,10 @@ use quickwit_proto::metastore::{
     MetastoreServiceClient, SourceType, ToggleSourceRequest,
 };
 use quickwit_proto::types::{IndexId, IndexUid, NodeId, ShardId, SourceId, SourceUid};
-pub(super) use shard_table::{ScalingMode, ShardEntry, ShardLocations, ShardStats, ShardTable};
+pub(super) use shard_table::{ShardEntry, ShardLocations, ShardStats, ShardTable};
 use tracing::{debug, error, info, instrument, warn};
 
+pub(super) use crate::ingest::ScalingMode;
 use crate::metrics::INDEXES_TOTAL;
 
 /// The control plane maintains a model in sync with the metastore.

quickwit/quickwit-control-plane/src/model/shard_table.rs

@@ -17,16 +17,18 @@ use std::collections::{BTreeSet, HashMap};
 use std::ops::{Deref, DerefMut};
 use std::time::Duration;
 
+use bytesize::ByteSize;
 use fnv::{FnvHashMap, FnvHashSet};
 use quickwit_common::metrics::index_label;
 use quickwit_common::rate_limiter::{RateLimiter, RateLimiterSettings};
 use quickwit_common::tower::ConstantRate;
-use quickwit_ingest::{RateMibPerSec, ShardInfo, ShardInfos};
+use quickwit_ingest::{ShardInfo, ShardInfos};
 use quickwit_metrics::{gauge, label_values};
 use quickwit_proto::ingest::{Shard, ShardState};
 use quickwit_proto::types::{IndexUid, NodeId, NodeIdRef, ShardId, SourceId, SourceUid};
 use tracing::{error, info, warn};
 
+use crate::ingest::ScalingMode;
 use crate::metrics::{CLOSED_SHARDS, INDEX_ID_LABEL_NAMES, OPEN_SHARDS};
 
 /// Limits the number of scale up operations that can happen to a source to 5 per minute.
@@ -43,19 +45,11 @@ const SCALING_DOWN_RATE_LIMITER_SETTINGS: RateLimiterSettings = RateLimiterSetti
     refill_period: Duration::from_mins(1),
 };
 
-#[derive(Debug, Clone, Copy, Eq, PartialEq)]
-pub(crate) enum ScalingMode {
-    /// Scale up by adding this number of shards
-    Up(usize),
-    /// Scale down by removing one shard
-    Down,
-}
-
 #[derive(Debug, Clone)]
 pub(crate) struct ShardEntry {
     pub shard: Shard,
-    pub short_term_ingestion_rate: RateMibPerSec,
-    pub long_term_ingestion_rate: RateMibPerSec,
+    pub short_term_ingestion_rate: ByteSize,
+    pub long_term_ingestion_rate: ByteSize,
 }
 
 impl Deref for ShardEntry {
@@ -76,8 +70,8 @@ impl From<Shard> for ShardEntry {
     fn from(shard: Shard) -> Self {
         Self {
             shard,
-            short_term_ingestion_rate: RateMibPerSec::default(),
-            long_term_ingestion_rate: RateMibPerSec::default(),
+            short_term_ingestion_rate: ByteSize::default(),
+            long_term_ingestion_rate: ByteSize::default(),
         }
     }
 }
@@ -109,25 +103,25 @@ impl ShardTableEntry {
     fn shards_stats(&self) -> ShardStats {
         let mut num_open_shards = 0;
         let mut num_closed_shards = 0;
-        let mut short_term_ingestion_rate_sum = 0;
-        let mut long_term_ingestion_rate_sum = 0;
+        let mut short_term_ingestion_rate_sum = ByteSize::default();
+        let mut long_term_ingestion_rate_sum = ByteSize::default();
 
         for shard_entry in self.shard_entries.values() {
             if shard_entry.is_open() {
                 num_open_shards += 1;
-                short_term_ingestion_rate_sum += shard_entry.short_term_ingestion_rate.0 as usize;
-                long_term_ingestion_rate_sum += shard_entry.long_term_ingestion_rate.0 as usize;
+                short_term_ingestion_rate_sum += shard_entry.short_term_ingestion_rate;
+                long_term_ingestion_rate_sum += shard_entry.long_term_ingestion_rate;
             } else if shard_entry.is_closed() {
                 num_closed_shards += 1;
             }
         }
         let avg_short_term_ingestion_rate = if num_open_shards > 0 {
-            short_term_ingestion_rate_sum as f32 / num_open_shards as f32
+            short_term_ingestion_rate_sum.as_mib() as f32 / num_open_shards as f32
         } else {
             0.0
         };
         let avg_long_term_ingestion_rate = if num_open_shards > 0 {
-            long_term_ingestion_rate_sum as f32 / num_open_shards as f32
+            long_term_ingestion_rate_sum.as_mib() as f32 / num_open_shards as f32
         } else {
             0.0
         };
@@ -895,32 +889,32 @@ mod tests {
             ShardInfo {
                 shard_id: ShardId::from(1),
                 shard_state: ShardState::Open,
-                short_term_ingestion_rate: RateMibPerSec(1),
-                long_term_ingestion_rate: RateMibPerSec(1),
+                short_term_ingestion_rate: ByteSize::mib(1),
+                long_term_ingestion_rate: ByteSize::mib(1),
             },
             ShardInfo {
                 shard_id: ShardId::from(2),
                 shard_state: ShardState::Open,
-                short_term_ingestion_rate: RateMibPerSec(2),
-                long_term_ingestion_rate: RateMibPerSec(2),
+                short_term_ingestion_rate: ByteSize::mib(2),
+                long_term_ingestion_rate: ByteSize::mib(2),
             },
             ShardInfo {
                 shard_id: ShardId::from(3),
                 shard_state: ShardState::Open,
-                short_term_ingestion_rate: RateMibPerSec(3),
-                long_term_ingestion_rate: RateMibPerSec(3),
+                short_term_ingestion_rate: ByteSize::mib(3),
+                long_term_ingestion_rate: ByteSize::mib(3),
             },
             ShardInfo {
                 shard_id: ShardId::from(4),
                 shard_state: ShardState::Closed,
-                short_term_ingestion_rate: RateMibPerSec(4),
-                long_term_ingestion_rate: RateMibPerSec(4),
+                short_term_ingestion_rate: ByteSize::mib(4),
+                long_term_ingestion_rate: ByteSize::mib(4),
             },
             ShardInfo {
                 shard_id: ShardId::from(5),
                 shard_state: ShardState::Open,
-                short_term_ingestion_rate: RateMibPerSec(5),
-                long_term_ingestion_rate: RateMibPerSec(5),
+                short_term_ingestion_rate: ByteSize::mib(5),
+                long_term_ingestion_rate: ByteSize::mib(5),
             },
         ]);
         let shard_stats = shard_table.update_shards(&source_uid, &shard_infos);
@@ -940,22 +934,22 @@ mod tests {
 
         assert_eq!(shard_entries[0].shard.shard_id(), ShardId::from(1));
         assert_eq!(shard_entries[0].shard.shard_state(), ShardState::Open);
-        assert_eq!(shard_entries[0].short_term_ingestion_rate, RateMibPerSec(1));
+        assert_eq!(shard_entries[0].short_term_ingestion_rate, ByteSize::mib(1));
 
         assert_eq!(shard_entries[1].shard.shard_id(), ShardId::from(2));
         assert_eq!(shard_entries[1].shard.shard_state(), ShardState::Open);
-        assert_eq!(shard_entries[1].short_term_ingestion_rate, RateMibPerSec(2));
+        assert_eq!(shard_entries[1].short_term_ingestion_rate, ByteSize::mib(2));
 
         assert_eq!(shard_entries[2].shard.shard_id(), ShardId::from(3));
         assert_eq!(
             shard_entries[2].shard.shard_state(),
             ShardState::Unavailable
         );
-        assert_eq!(shard_entries[2].short_term_ingestion_rate, RateMibPerSec(3));
+        assert_eq!(shard_entries[2].short_term_ingestion_rate, ByteSize::mib(3));
 
         assert_eq!(shard_entries[3].shard.shard_id(), ShardId::from(4));
         assert_eq!(shard_entries[3].shard.shard_state(), ShardState::Closed);
-        assert_eq!(shard_entries[3].short_term_ingestion_rate, RateMibPerSec(4));
+        assert_eq!(shard_entries[3].short_term_ingestion_rate, ByteSize::mib(4));
     }
 
     #[test]

quickwit/quickwit-ingest/src/ingest_v2/broadcast/local_shards.rs

@@ -33,18 +33,16 @@ use crate::RateMibPerSec;
 use crate::ingest_v2::shard_readings::{ShardReadingsBySource, ShardThroughputReading};
 use crate::ingest_v2::state::WeakIngesterState;
 
-const ONE_MIB: ByteSize = ByteSize::mib(1);
-
 /// Broadcasted information about a shard.
 #[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
 pub struct ShardInfo {
     pub shard_id: ShardId,
     pub shard_state: ShardState,
-    /// Shard ingestion rate in MiB/s.
+    /// Shard ingestion rate in bytes/s, serialized as rounded-up MiB/s for gossip.
     /// Short term ingestion rate. It is measured over a short period of time.
-    pub short_term_ingestion_rate: RateMibPerSec,
+    pub short_term_ingestion_rate: ByteSize,
     /// Long term ingestion rate. It is measured over a larger period of time.
-    pub long_term_ingestion_rate: RateMibPerSec,
+    pub long_term_ingestion_rate: ByteSize,
 }
 
 impl Serialize for ShardInfo {
@@ -53,8 +51,8 @@ impl Serialize for ShardInfo {
             "{}:{}:{}:{}",
             self.shard_id,
             self.shard_state.as_json_str_name(),
-            self.short_term_ingestion_rate.0,
-            self.long_term_ingestion_rate.0,
+            RateMibPerSec::from(self.short_term_ingestion_rate).0,
+            RateMibPerSec::from(self.long_term_ingestion_rate).0,
         ))
     }
 }
@@ -81,6 +79,7 @@ impl<'de> Deserialize<'de> for ShardInfo {
             .ok_or_else(|| serde::de::Error::custom("invalid shard info"))?
             .parse::<u16>()
             .map(RateMibPerSec)
+            .map(ByteSize::from)
             .map_err(|_| serde::de::Error::custom("invalid shard ingestion rate"))?;
 
         let long_term_ingestion_rate = parts
@@ -88,6 +87,7 @@ impl<'de> Deserialize<'de> for ShardInfo {
             .ok_or_else(|| serde::de::Error::custom("invalid shard info"))?
             .parse::<u16>()
             .map(RateMibPerSec)
+            .map(ByteSize::from)
             .map_err(|_| serde::de::Error::custom("invalid shard ingestion rate"))?;
 
         Ok(Self {
@@ -101,23 +101,11 @@ impl<'de> Deserialize<'de> for ShardInfo {
 
 impl From<&ShardThroughputReading> for ShardInfo {
     fn from(reading: &ShardThroughputReading) -> Self {
-        let short_term_ingestion_rate_mib_per_sec_u64: u64 = reading
-            .short_term_ingestion_rate
-            .as_u64()
-            .div_ceil(ONE_MIB.as_u64());
-        let long_term_ingestion_rate_mib_per_sec_u64: u64 = reading
-            .long_term_ingestion_rate
-            .as_u64()
-            .div_ceil(ONE_MIB.as_u64());
         Self {
             shard_id: reading.shard_id.clone(),
             shard_state: reading.shard_state,
-            short_term_ingestion_rate: RateMibPerSec(
-                short_term_ingestion_rate_mib_per_sec_u64 as u16,
-            ),
-            long_term_ingestion_rate: RateMibPerSec(
-                long_term_ingestion_rate_mib_per_sec_u64 as u16,
-            ),
+            short_term_ingestion_rate: reading.short_term_ingestion_rate,
+            long_term_ingestion_rate: reading.long_term_ingestion_rate,
         }
     }
 }
@@ -322,7 +310,6 @@ mod tests {
     use quickwit_proto::types::{IndexUid, ShardId, SourceId, SourceUid};
 
     use super::*;
-    use crate::RateMibPerSec;
     use crate::ingest_v2::state::IngesterState;
 
     #[test]
@@ -330,8 +317,8 @@ mod tests {
         let shard_info = ShardInfo {
             shard_id: ShardId::from(1),
             shard_state: ShardState::Open,
-            short_term_ingestion_rate: RateMibPerSec(42),
-            long_term_ingestion_rate: RateMibPerSec(40),
+            short_term_ingestion_rate: ByteSize::mib(42),
+            long_term_ingestion_rate: ByteSize::mib(40),
         };
         let serialized = serde_json::to_string(&shard_info).unwrap();
         assert_eq!(serialized, r#""00000000000000000001:open:42:40""#);
@@ -358,8 +345,8 @@ mod tests {
                 vec![ShardInfo {
                     shard_id: ShardId::from(1),
                     shard_state: ShardState::Open,
-                    short_term_ingestion_rate: RateMibPerSec(42),
-                    long_term_ingestion_rate: RateMibPerSec(42),
+                    short_term_ingestion_rate: ByteSize::mib(42),
+                    long_term_ingestion_rate: ByteSize::mib(42),
                 }]
                 .into_iter()
                 .collect(),
@@ -399,8 +386,8 @@ mod tests {
                 vec![ShardInfo {
                     shard_id: ShardId::from(1),
                     shard_state: ShardState::Closed,
-                    short_term_ingestion_rate: RateMibPerSec(42),
-                    long_term_ingestion_rate: RateMibPerSec(42),
+                    short_term_ingestion_rate: ByteSize::mib(42),
+                    long_term_ingestion_rate: ByteSize::mib(42),
                 }]
                 .into_iter()
                 .collect(),
@@ -539,7 +526,7 @@ mod tests {
                 let shard_info = event.shard_infos.iter().next().unwrap();
                 assert_eq!(shard_info.shard_id, ShardId::from(1));
                 assert_eq!(shard_info.shard_state, ShardState::Open);
-                assert_eq!(shard_info.short_term_ingestion_rate, 42u16);
+                assert_eq!(shard_info.short_term_ingestion_rate, ByteSize::mib(42));
             })
             .forever();
 
@@ -555,8 +542,8 @@ mod tests {
         let value = serde_json::to_string(&vec![ShardInfo {
             shard_id: ShardId::from(1),
             shard_state: ShardState::Open,
-            short_term_ingestion_rate: RateMibPerSec(42),
-            long_term_ingestion_rate: RateMibPerSec(42),
+            short_term_ingestion_rate: ByteSize::mib(42),
+            long_term_ingestion_rate: ByteSize::mib(42),
         }])
         .unwrap();
 

quickwit/quickwit-ingest/src/ingest_v2/ingester.rs

@@ -1437,7 +1437,7 @@ mod tests {
         let shard_info = shard_infos.iter().next().unwrap();
         assert_eq!(shard_info.shard_id, ShardId::from(1));
         assert_eq!(shard_info.shard_state, ShardState::Open);
-        assert_eq!(shard_info.short_term_ingestion_rate, 0);
+        assert_eq!(shard_info.short_term_ingestion_rate, ByteSize::b(0));
 
         let mut state_guard = ingester.state.lock_fully("test").await.unwrap();
         state_guard.shards.get_mut(&queue_id_01).unwrap().close();

quickwit/quickwit-ingest/src/ingest_v2/mod.rs

@@ -310,6 +310,18 @@ impl AddAssign<RateMibPerSec> for RateMibPerSec {
     }
 }
 
+impl From<ByteSize> for RateMibPerSec {
+    fn from(rate: ByteSize) -> Self {
+        RateMibPerSec(rate.as_u64().div_ceil(ByteSize::mib(1).as_u64()) as u16)
+    }
+}
+
+impl From<RateMibPerSec> for ByteSize {
+    fn from(rate: RateMibPerSec) -> Self {
+        ByteSize::mib(rate.0.into())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use bytes::Bytes;