Optimize rocksdb

src/cli/sync.rs

@@ -6,6 +6,7 @@ use colored::Colorize;
 use crate::config::types::ApplicationConfig;
 use crate::constants::{DB_CF_OPTIONS, DB_OPTIONS};
 use crate::crawler::{dlsite, DLSiteCrawler};
+use crate::helpers;
 use crate::helpers::db::RocksDB;
 use crate::models::DLSiteManiax;
 
@@ -63,29 +64,23 @@ impl SyncDLSiteCommand {
     async fn sync_works(app_conf: &ApplicationConfig, db: &RocksDB) -> Result<()> {
         let crawler = DLSiteCrawler::new();
         let mut rj_nums: Vec<String> = Vec::new();
-        for path_str in app_conf.path_config.dlsite_paths.iter() {
-            let path = Path::new(path_str);
-            if !path.exists() {
-                return Err(eyre!("{} {}", path_str.blue(), "does not exist".red()));
+        let paths = app_conf.path_config.dlsite_paths.iter()
+            .map(|path| Path::new(path).to_path_buf())
+            .collect::<Vec<_>>();
+        let dirs = helpers::get_all_folders(&paths).await?;
+        for dir_path in dirs.iter() {
+            if !dir_path.is_dir() {
+                println!("{dir_path:?} is not a directory");
+                continue;
             }
-            let dir_paths = path.read_dir()?
-                .filter_map(Result::ok)
-                .map(|e| e.path())
-                .collect::<Vec<_>>();
-            for dir_path in dir_paths.iter() {
-                if !dir_path.is_dir() {
-                    println!("{dir_path:?} is not a directory");
-                    continue;
-                }
-                let dir_name = dir_path
-                    .file_name().unwrap()
-                    .to_str().unwrap();
-                if !dlsite::is_valid_rj_number(dir_name) {
-                    println!("{} {}", dir_path.to_str().unwrap().blue(), "is not a valid rj number, please add it manually".red());
-                    continue;
-                }
-                rj_nums.push(dir_name.to_string());
+            let dir_name = dir_path
+                .file_name().unwrap()
+                .to_str().unwrap();
+            if !dlsite::is_valid_rj_number(dir_name) {
+                println!("{} {}", dir_path.to_str().unwrap().blue(), "is not a valid rj number, please add it manually".red());
+                continue;
             }
+            rj_nums.push(dir_name.to_string());
         }
         let maniaxes = crawler.get_game_infos(rj_nums).await?;
         db.set_values(&maniaxes)?;

src/constants.rs

@@ -14,7 +14,7 @@ lazy_static! {
     pub static ref APP_DB_DATA_DIR: PathBuf = APP_DATA_DIR.clone().join("db");
 
     pub static ref DB_OPTIONS: rocksdb::Options = get_db_options();
-    pub static ref DB_CF_OPTIONS: rocksdb::Options = get_cf_options();
+    pub static ref DB_CF_OPTIONS: rocksdb::Options = rocksdb::Options::default();
 }
 
 lazy_static! {
@@ -26,12 +26,7 @@ fn get_db_options() -> rocksdb::Options {
 
     opts.create_missing_column_families(true);
     opts.create_if_missing(true);
-
-    opts
-}
-
-fn get_cf_options() -> rocksdb::Options {
-    let opts = rocksdb::Options::default();
+    opts.increase_parallelism(num_cpus::get() as i32);
 
     opts
 }

src/helpers/db.rs

@@ -1,5 +1,5 @@
 use crate::constants::{APP_DB_DATA_DIR, DB_COLUMNS};
-use rocksdb::{ColumnFamilyDescriptor, IteratorMode, OptimisticTransactionDB, Options};
+use rocksdb::{ColumnFamilyDescriptor, IteratorMode, OptimisticTransactionDB, Options, ReadOptions};
 use serde::{Serialize};
 use serde::de::DeserializeOwned;
 use crate::models::RocksColumn;
@@ -63,8 +63,10 @@ impl RocksDB {
         where TColumn: RocksColumn + DeserializeOwned
     {
         let cf = self.db.cf_handle(TColumn::get_column_name().as_str()).unwrap();
-        let values = self.db.iterator_cf(&cf, IteratorMode::Start)
-            .filter_map(|res| res.ok())
+        let mut options = ReadOptions::default();
+        options.set_async_io(true);
+        let values = self.db.iterator_cf_opt(&cf, options, IteratorMode::Start)
+            .filter_map(Result::ok)
             .map(|(k, v)|
                 (
                     serde_json::from_slice::<TColumn::Id>(&k).unwrap(),

src/helpers/mod.rs

@@ -1,5 +1,8 @@
 pub mod db;
 
+use std::path::PathBuf;
+use color_eyre::eyre::eyre;
+use color_eyre::owo_colors::OwoColorize;
 use tokio::fs;
 use crate::constants::{APP_CONFIG_DIR, APP_DATA_DIR, APP_DB_DATA_DIR};
 use crate::crawler::DLSITE_IMG_FOLDER;
@@ -19,4 +22,20 @@ pub async fn initialize_folders() -> color_eyre::Result<()> {
         fs::create_dir_all(APP_DB_DATA_DIR.as_path()).await?;
     }
     Ok(())
+}
+
+pub async fn get_all_folders(paths: &Vec<PathBuf>) -> color_eyre::Result<Vec<PathBuf>> {
+    let mut folders: Vec<PathBuf> = Vec::new();
+    for path in paths {
+        let path = path.as_path();
+        if !path.exists() {
+            return Err(eyre!("{:?} {}", path.blue(), "does not exist".red()));
+        }
+
+        let mut dirs = fs::read_dir(path).await?;
+        while let Some(dir) = dirs.next_entry().await? {
+            folders.push(dir.path());
+        }
+    }
+    Ok(folders)
 }