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

meilisearch-http/src/analytics.rs

@@ -8,6 +8,7 @@ use serde_qs as qs;
 use siphasher::sip::SipHasher;
 use walkdir::WalkDir;
 
+use crate::helpers::EnvSizer;
 use crate::Data;
 use crate::Opt;
 
@@ -33,12 +34,7 @@ impl EventProperties {
             }
         }
 
-        let database_size = WalkDir::new(&data.db_path)
-            .into_iter()
-            .filter_map(|entry| entry.ok())
-            .filter_map(|entry| entry.metadata().ok())
-            .filter(|metadata| metadata.is_file())
-            .fold(0, |acc, m| acc + m.len());
+        let database_size = data.env.size();
 
         let last_update_timestamp = data.db.last_update(&reader)?.map(|u| u.timestamp());
 
@@ -116,7 +112,7 @@ pub fn analytics_sender(data: Data, opt: Opt) {
             time,
             app_version,
             user_properties,
-            event_properties
+            event_properties,
         };
         let event = serde_json::to_string(&event).unwrap();
 

meilisearch-http/src/data/mod.rs

@@ -1,14 +1,14 @@
+use std::collections::HashMap;
 use std::ops::Deref;
 use std::sync::Arc;
 
+use chrono::{DateTime, Utc};
 use sha2::Digest;
 
 use crate::index::Settings;
 use crate::index_controller::{IndexController, IndexStats};
 use crate::index_controller::{IndexMetadata, IndexSettings};
 use crate::option::Opt;
-use std::collections::HashMap;
-use chrono::{DateTime, Utc};
 
 pub mod search;
 mod updates;
@@ -119,15 +119,22 @@ impl Data {
 
     pub async fn get_stats(&self) -> anyhow::Result<Stats> {
         let mut stats = Stats::default();
+        stats.database_size += self.index_controller.get_uuids_size().await?;
 
         for index in self.index_controller.list_indexes().await? {
             let index_stats = self.index_controller.get_stats(index.uid.clone()).await?;
 
             stats.database_size += index_stats.size;
+            stats.database_size += self
+                .index_controller
+                .get_updates_size(index.uid.clone())
+                .await?;
+
             stats.last_update = Some(match stats.last_update {
                 Some(last_update) => last_update.max(index.meta.updated_at),
                 None => index.meta.updated_at,
             });
+
             stats.indexes.insert(index.uid, index_stats);
         }
 

meilisearch-http/src/helpers/env.rs

@@ -0,0 +1,16 @@
+use walkdir::WalkDir;
+
+pub trait EnvSizer {
+    fn size(&self) -> u64;
+}
+
+impl EnvSizer for heed::Env {
+    fn size(&self) -> u64 {
+        WalkDir::new(self.path())
+            .into_iter()
+            .filter_map(|entry| entry.ok())
+            .filter_map(|entry| entry.metadata().ok())
+            .filter(|metadata| metadata.is_file())
+            .fold(0, |acc, m| acc + m.len())
+    }
+}

meilisearch-http/src/helpers/mod.rs

@@ -1,4 +1,6 @@
 pub mod authentication;
 pub mod compression;
+mod env;
 
 pub use authentication::Authentication;
+pub use env::EnvSizer;

meilisearch-http/src/index/mod.rs

@@ -5,10 +5,10 @@ use std::sync::Arc;
 use anyhow::{bail, Context};
 use milli::obkv_to_json;
 use serde_json::{Map, Value};
-use walkdir::WalkDir;
 
 pub use search::{SearchQuery, SearchResult, DEFAULT_SEARCH_LIMIT};
 pub use updates::{Facets, Settings, UpdateResult};
+use crate::helpers::EnvSizer;
 
 mod search;
 mod updates;
@@ -55,11 +55,7 @@ impl Index {
         let stop_words = self
             .stop_words(&txn)?
             .map(|stop_words| -> anyhow::Result<BTreeSet<_>> {
-                Ok(stop_words
-                .stream()
-                .into_strs()?
-                .into_iter()
-                .collect())
+                Ok(stop_words.stream().into_strs()?.into_iter().collect())
             })
             .transpose()?
             .unwrap_or_else(BTreeSet::new);
@@ -127,13 +123,8 @@ impl Index {
         }
     }
 
-    pub fn size(&self) -> anyhow::Result<u64> {
-        Ok(WalkDir::new(self.env.path())
-            .into_iter()
-            .filter_map(|entry| entry.ok())
-            .filter_map(|entry| entry.metadata().ok())
-            .filter(|metadata| metadata.is_file())
-            .fold(0, |acc, m| acc + m.len()))
+    pub fn size(&self) -> u64 {
+        self.env.size()
     }
 
     fn fields_to_display<S: AsRef<str>>(

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())
+    }
 }