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

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