From 350d4eb5b0b93b153c3e94755720fcae913ea031 Mon Sep 17 00:00:00 2001
From: Louis Dureuil <louis@meilisearch.com>
Date: Tue, 28 Jan 2025 11:08:05 +0100
Subject: [PATCH] Network stored in DB

---
 crates/index-scheduler/src/features.rs    | 32 ++++++++++++++++++++---
 crates/index-scheduler/src/lib.rs         | 13 +++++++--
 crates/meilisearch/src/routes/features.rs |  2 +-
 3 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/crates/index-scheduler/src/features.rs b/crates/index-scheduler/src/features.rs
index c6c17b2d5..143b96857 100644
--- a/crates/index-scheduler/src/features.rs
+++ b/crates/index-scheduler/src/features.rs
@@ -1,6 +1,6 @@
 use std::sync::{Arc, RwLock};
 
-use meilisearch_types::features::{InstanceTogglableFeatures, RuntimeTogglableFeatures};
+use meilisearch_types::features::{InstanceTogglableFeatures, Network, RuntimeTogglableFeatures};
 use meilisearch_types::heed::types::{SerdeJson, Str};
 use meilisearch_types::heed::{Database, Env, RwTxn};
 
@@ -8,16 +8,18 @@ use crate::error::FeatureNotEnabledError;
 use crate::Result;
 
 /// The number of database used by features
-const NUMBER_OF_DATABASES: u32 = 1;
+const NUMBER_OF_DATABASES: u32 = 2;
 /// Database const names for the `FeatureData`.
 mod db_name {
     pub const EXPERIMENTAL_FEATURES: &str = "experimental-features";
+    const NETWORK: &str = "network";
 }
 
 #[derive(Clone)]
 pub(crate) struct FeatureData {
     persisted: Database<Str, SerdeJson<RuntimeTogglableFeatures>>,
     runtime: Arc<RwLock<RuntimeTogglableFeatures>>,
+    network: Arc<RwLock<Network>>,
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -111,7 +113,14 @@ impl FeatureData {
             ..persisted_features
         }));
 
-        Ok(Self { persisted: runtime_features_db, runtime })
+        let network_db = runtime_features_db.remap_data_type::<SerdeJson<Network>>();
+        let network: Network = network_db.get(&txn, db_name::NETWORK)?.unwrap_or_default();
+
+        Ok(Self {
+            persisted: runtime_features_db,
+            runtime,
+            network: Arc::new(RwLock::new(network)),
+        })
     }
 
     pub fn put_runtime_features(
@@ -140,4 +149,21 @@ impl FeatureData {
     pub fn features(&self) -> RoFeatures {
         RoFeatures::new(self)
     }
+
+    pub fn put_network(&self, mut wtxn: RwTxn, new_network: Network) -> Result<()> {
+        self.persisted.remap_data_type::<SerdeJson<Network>>().put(
+            &mut wtxn,
+            db_name::NETWORK,
+            &new_network,
+        )?;
+        wtxn.commit()?;
+
+        let mut network = self.network.write().unwrap();
+        *network = new_network;
+        Ok(())
+    }
+
+    pub fn network(&self) -> Network {
+        Network::clone(&*self.network.read().unwrap())
+    }
 }
diff --git a/crates/index-scheduler/src/lib.rs b/crates/index-scheduler/src/lib.rs
index 530b7bedc..0f8212470 100644
--- a/crates/index-scheduler/src/lib.rs
+++ b/crates/index-scheduler/src/lib.rs
@@ -51,7 +51,7 @@ pub use features::RoFeatures;
 use flate2::bufread::GzEncoder;
 use flate2::Compression;
 use meilisearch_types::batches::Batch;
-use meilisearch_types::features::{InstanceTogglableFeatures, RuntimeTogglableFeatures};
+use meilisearch_types::features::{InstanceTogglableFeatures, Network, RuntimeTogglableFeatures};
 use meilisearch_types::heed::byteorder::BE;
 use meilisearch_types::heed::types::I128;
 use meilisearch_types::heed::{self, Env, RoTxn};
@@ -770,7 +770,16 @@ impl IndexScheduler {
         Ok(())
     }
 
-    // TODO: consider using a type alias or a struct embedder/template
+    pub fn put_network(&self, network: Network) -> Result<()> {
+        let wtxn = self.env.write_txn().map_err(Error::HeedTransaction)?;
+        self.features.put_network(wtxn, network)?;
+        Ok(())
+    }
+
+    pub fn network(&self) -> Network {
+        self.features.network()
+    }
+
     pub fn embedders(
         &self,
         index_uid: String,
diff --git a/crates/meilisearch/src/routes/features.rs b/crates/meilisearch/src/routes/features.rs
index f46bda5a0..e57cbcaf0 100644
--- a/crates/meilisearch/src/routes/features.rs
+++ b/crates/meilisearch/src/routes/features.rs
@@ -184,7 +184,7 @@ async fn patch_features(
     };
 
     // explicitly destructure for analytics rather than using the `Serialize` implementation, because
-    // the it renames to camelCase, which we don't want for analytics.
+    // it renames to camelCase, which we don't want for analytics.
     // **Do not** ignore fields with `..` or `_` here, because we want to add them in the future.
     let meilisearch_types::features::RuntimeTogglableFeatures {
         metrics,