perf: cache for players scraping

sdgb-cli/Cargo.toml

@@ -6,12 +6,15 @@ authors = ["mokurin000"]
 description = "CLI tool for SDGB protocol"
 
 [features]
-default = ["compio"]
+default = ["compio", "cache"]
+
 compio = ["dep:compio", "sdgb-api/compio"]
 tokio = ["dep:tokio", "sdgb-api/tokio"]
 
+cache = ["dep:redb"]
+
 [dependencies]
-sdgb-api = { workspace = true }
+sdgb-api = { workspace = true, features = ["bincode"] }
 
 snafu = { workspace = true }
 serde_json = { workspace = true }
@@ -27,3 +30,4 @@ spdlog-rs = { version = "0.4.3", default-features = false, features = [
     "release-level-info",
 ] }
 futures-util = "0.3.31"
+redb = { version = "2.6.1", optional = true }

sdgb-cli/src/cache/mod.rs

@@ -0,0 +1,25 @@
+use std::sync::LazyLock;
+
+use redb::{Table, TableDefinition, WriteTransaction};
+
+static DATABASE: LazyLock<redb::Database> = LazyLock::new(|| {
+    redb::Database::builder()
+        .create("players.redb")
+        .expect("failed to open database")
+});
+const DIFINITION: TableDefinition<'_, u32, Vec<u8>> = redb::TableDefinition::new("players");
+
+pub fn write_txn() -> Result<WriteTransaction, redb::Error> {
+    Ok(DATABASE.begin_write()?)
+}
+
+pub fn open_table(write: &WriteTransaction) -> Result<Table<'_, u32, Vec<u8>>, redb::Error> {
+    Ok(write.open_table(DIFINITION)?)
+}
+
+pub fn init_db() -> Result<(), redb::Error> {
+    let write_txn = DATABASE.begin_write()?;
+    write_txn.open_table(DIFINITION)?;
+    write_txn.commit()?;
+    Ok(())
+}

sdgb-cli/src/main.rs

@@ -6,6 +6,7 @@ use std::{
 use futures_util::StreamExt;
 use nyquest_preset::nyquest::ClientBuilder;
 use palc::Parser;
+
 use sdgb_api::{
     ApiError,
     all_net::QRCode,
@@ -23,6 +24,8 @@ use spdlog::{error, info, warn};
 
 use crate::commands::Cli;
 
+#[cfg(feature = "cache")]
+mod cache;
 mod commands;
 
 #[cfg_attr(feature = "compio", compio::main)]
@@ -102,9 +105,33 @@ async fn main() -> Result<(), Box<dyn snafu::Error>> {
                 user_ids.push(user_id);
             }
 
-            let users = futures_util::stream::iter(user_ids)
+            #[cfg(feature = "cache")]
+            let _ = cache::init_db();
+            #[cfg(feature = "cache")]
+            let write = cache::write_txn()?;
+            #[cfg(feature = "cache")]
+            let config = sdgb_api::bincode::config::Configuration::<
+                sdgb_api::bincode::config::LittleEndian,
+            >::default()
+            .with_no_limit();
+
+            let players = futures_util::stream::iter(user_ids)
                 .map(async |user_id| {
-                    info!("preview: {user_id}");
+                    #[cfg(feature = "cache")]
+                    {
+                        use redb::ReadableTable;
+                        use sdgb_api::bincode::borrow_decode_from_slice;
+
+                        let cache_table = cache::open_table(&write)?;
+                        let data = cache_table.get(user_id)?;
+                        if let Some(data) = data {
+                            let decoded: (GetUserPreviewApiResp, _) =
+                                borrow_decode_from_slice(&data.value(), config)?;
+
+                            return Ok(decoded.0);
+                        }
+                    }
+
                     let resp = Sdgb1_50::request::<_, GetUserPreviewApiResp>(
                         &client,
                         APIMethod::GetUserPreviewApi,
@@ -113,15 +140,30 @@ async fn main() -> Result<(), Box<dyn snafu::Error>> {
                     )
                     .await;
 
-                    if let Err(e) = &resp {
-                        if matches!(e, ApiError::JSON { .. }) {
+                    match &resp {
+                        Ok(resp) => {
+                            info!("preview: {user_id} succeed");
+
+                            #[cfg(feature = "cache")]
+                            {
+                                use sdgb_api::bincode::encode_to_vec;
+
+                                if let Ok(mut table) = cache::open_table(&write)
+                                    && let Ok(encoded) = encode_to_vec(resp, config)
+                                {
+                                    _ = table.insert(resp.user_id, encoded);
+                                }
+                            }
+                        }
+                        Err(ApiError::JSON { .. }) => {
                             warn!("account unregistered: {user_id}");
-                        } else {
+                        }
+                        Err(e) => {
                             error!("preview failed: {e}");
                         }
                     }
 
-                    resp
+                    Result::<_, Box<dyn snafu::Error>>::Ok(resp?)
                 })
                 .buffer_unordered(20)
                 .filter_map(async |r| r.ok())
@@ -132,8 +174,8 @@ async fn main() -> Result<(), Box<dyn snafu::Error>> {
                 .write(true)
                 .truncate(true)
                 .create(true)
-                .open("users.json")?;
-            serde_json::to_writer_pretty(output, &users)?;
+                .open("players.json")?;
+            serde_json::to_writer_pretty(output, &players)?;
         }
     }