feat(http): calculate updates' and uuids' dbs size

meilisearch-http/src/index_controller/index_actor/actor.rs

@@ -360,7 +360,7 @@ impl<S: IndexStore + Sync + Send> IndexActor<S> {
             let rtxn = index.read_txn()?;
 
             Ok(IndexStats {
-                size: index.size()?,
+                size: index.size(),
                 number_of_documents: index.number_of_documents(&rtxn)?,
                 is_indexing,
                 fields_distribution: index.fields_distribution(&rtxn)?,

meilisearch-http/src/index_controller/mod.rs

@@ -356,6 +356,16 @@ impl IndexController {
 
         Ok(self.index_handle.get_index_stats(uuid).await?)
     }
+
+    pub async fn get_updates_size(&self, uid: String) -> anyhow::Result<u64> {
+        let uuid = self.uuid_resolver.get(uid.clone()).await?;
+
+        Ok(self.update_handle.get_size(uuid).await?)
+    }
+
+    pub async fn get_uuids_size(&self) -> anyhow::Result<u64> {
+        Ok(self.uuid_resolver.get_size().await?)
+    }
 }
 
 pub async fn get_arc_ownership_blocking<T>(mut item: Arc<T>) -> T {

meilisearch-http/src/index_controller/update_actor/actor.rs

@@ -8,10 +8,11 @@ use tokio::io::{AsyncSeekExt, AsyncWriteExt};
 use tokio::sync::mpsc;
 use uuid::Uuid;
 
-use super::{PayloadData, Result, UpdateError, UpdateMsg, UpdateStoreStore};
 use crate::index_controller::index_actor::IndexActorHandle;
 use crate::index_controller::{get_arc_ownership_blocking, UpdateMeta, UpdateStatus};
 
+use super::{PayloadData, Result, UpdateError, UpdateMsg, UpdateStoreStore};
+
 pub struct UpdateActor<D, S, I> {
     path: PathBuf,
     store: S,
@@ -72,6 +73,9 @@ where
                 Some(Snapshot { uuid, path, ret }) => {
                     let _ = ret.send(self.handle_snapshot(uuid, path).await);
                 }
+                Some(GetSize { uuid, ret }) => {
+                    let _ = ret.send(self.handle_get_size(uuid).await);
+                }
                 None => break,
             }
         }
@@ -223,4 +227,20 @@ where
 
         Ok(())
     }
+
+    async fn handle_get_size(&self, uuid: Uuid) -> Result<u64> {
+        let size = match self.store.get(uuid).await? {
+            Some(update_store) => tokio::task::spawn_blocking(move || -> anyhow::Result<u64> {
+                let txn = update_store.env.read_txn()?;
+
+                update_store.get_size(&txn)
+            })
+            .await
+            .map_err(|e| UpdateError::Error(e.into()))?
+            .map_err(|e| UpdateError::Error(e.into()))?,
+            None => 0,
+        };
+
+        Ok(size)
+    }
 }

meilisearch-http/src/index_controller/update_actor/handle_impl.rs

@@ -79,6 +79,13 @@ where
         receiver.await.expect("update actor killed.")
     }
 
+    async fn get_size(&self, uuid: Uuid) -> Result<u64> {
+        let (ret, receiver) = oneshot::channel();
+        let msg = UpdateMsg::GetSize { uuid, ret };
+        let _ = self.sender.send(msg).await;
+        receiver.await.expect("update actor killed.")
+    }
+
     async fn update(
         &self,
         meta: UpdateMeta,

meilisearch-http/src/index_controller/update_actor/message.rs

@@ -34,4 +34,8 @@ pub enum UpdateMsg<D> {
         path: PathBuf,
         ret: oneshot::Sender<Result<()>>,
     },
+    GetSize {
+        uuid: Uuid,
+        ret: oneshot::Sender<Result<u64>>,
+    },
 }

meilisearch-http/src/index_controller/update_actor/mod.rs

@@ -46,6 +46,7 @@ pub trait UpdateActorHandle {
     async fn delete(&self, uuid: Uuid) -> Result<()>;
     async fn create(&self, uuid: Uuid) -> Result<()>;
     async fn snapshot(&self, uuid: Uuid, path: PathBuf) -> Result<()>;
+    async fn get_size(&self, uuid: Uuid) -> Result<u64>;
     async fn update(
         &self,
         meta: UpdateMeta,

meilisearch-http/src/index_controller/update_actor/update_store.rs

@@ -1,3 +1,4 @@
+use std::fs::File;
 use std::fs::{copy, create_dir_all, remove_file};
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
@@ -6,10 +7,10 @@ use heed::types::{DecodeIgnore, OwnedType, SerdeJson};
 use heed::{CompactionOption, Database, Env, EnvOpenOptions};
 use parking_lot::{Mutex, RwLock};
 use serde::{Deserialize, Serialize};
-use std::fs::File;
 use tokio::sync::mpsc;
 use uuid::Uuid;
 
+use crate::helpers::EnvSizer;
 use crate::index_controller::updates::*;
 
 type BEU64 = heed::zerocopy::U64<heed::byteorder::BE>;
@@ -409,4 +410,18 @@ where
 
         Ok(())
     }
+
+    pub fn get_size(&self, txn: &heed::RoTxn) -> anyhow::Result<u64> {
+        let mut size = self.env.size();
+
+        for path in self.pending.iter(txn)? {
+            let (_, path) = path?;
+
+            if let Ok(metadata) = path.metadata() {
+                size += metadata.len()
+            }
+        }
+
+        Ok(size)
+    }
 }

meilisearch-http/src/index_controller/uuid_resolver/actor.rs

@@ -41,6 +41,9 @@ impl<S: UuidStore> UuidResolverActor<S> {
                 Some(SnapshotRequest { path, ret }) => {
                     let _ = ret.send(self.handle_snapshot(path).await);
                 }
+                Some(GetSize { ret }) => {
+                    let _ = ret.send(self.handle_get_size().await);
+                }
                 // all senders have been dropped, need to quit.
                 None => break,
             }
@@ -86,6 +89,10 @@ impl<S: UuidStore> UuidResolverActor<S> {
         self.store.insert(uid, uuid).await?;
         Ok(())
     }
+
+    async fn handle_get_size(&self) -> Result<u64> {
+        self.store.get_size().await
+    }
 }
 
 fn is_index_uid_valid(uid: &str) -> bool {

meilisearch-http/src/index_controller/uuid_resolver/handle_impl.rs

@@ -75,4 +75,13 @@ impl UuidResolverHandle for UuidResolverHandleImpl {
             .await
             .expect("Uuid resolver actor has been killed")?)
     }
+
+    async fn get_size(&self) -> Result<u64> {
+        let (ret, receiver) = oneshot::channel();
+        let msg = UuidResolveMsg::GetSize { ret };
+        let _ = self.sender.send(msg).await;
+        Ok(receiver
+            .await
+            .expect("Uuid resolver actor has been killed")?)
+    }
 }

meilisearch-http/src/index_controller/uuid_resolver/message.rs

@@ -4,6 +4,7 @@ use tokio::sync::oneshot;
 use uuid::Uuid;
 
 use super::Result;
+
 pub enum UuidResolveMsg {
     Get {
         uid: String,
@@ -29,4 +30,7 @@ pub enum UuidResolveMsg {
         path: PathBuf,
         ret: oneshot::Sender<Result<Vec<Uuid>>>,
     },
+    GetSize {
+        ret: oneshot::Sender<Result<u64>>,
+    },
 }

meilisearch-http/src/index_controller/uuid_resolver/mod.rs

@@ -30,6 +30,7 @@ pub trait UuidResolverHandle {
     async fn delete(&self, name: String) -> anyhow::Result<Uuid>;
     async fn list(&self) -> anyhow::Result<Vec<(String, Uuid)>>;
     async fn snapshot(&self, path: PathBuf) -> Result<Vec<Uuid>>;
+    async fn get_size(&self) -> Result<u64>;
 }
 
 #[derive(Debug, Error)]

meilisearch-http/src/index_controller/uuid_resolver/store.rs

@@ -8,6 +8,7 @@ use heed::{
 use uuid::Uuid;
 
 use super::{Result, UuidError, UUID_STORE_SIZE};
+use crate::helpers::EnvSizer;
 
 #[async_trait::async_trait]
 pub trait UuidStore {
@@ -19,6 +20,7 @@ pub trait UuidStore {
     async fn list(&self) -> Result<Vec<(String, Uuid)>>;
     async fn insert(&self, name: String, uuid: Uuid) -> Result<()>;
     async fn snapshot(&self, path: PathBuf) -> Result<Vec<Uuid>>;
+    async fn get_size(&self) -> Result<u64>;
 }
 
 pub struct HeedUuidStore {
@@ -151,4 +153,8 @@ impl UuidStore for HeedUuidStore {
         })
         .await?
     }
+
+    async fn get_size(&self) -> Result<u64> {
+        Ok(self.env.size())
+    }
 }