Fix friend issue

Cargo.toml

@@ -10,7 +10,7 @@ members = [
 resolver = "2"
 
 [workspace.package]
-version = "0.1.5"
+version = "0.1.6"
 
 [workspace.dependencies]
 anyhow = "1.0.102"

cuscuta-common/src/db/log/event.rs

@@ -1,5 +1,5 @@
 use chrono::Utc;
-use redis::Client;
+use redis::{Client, TypedCommands};
 
 use crate::db::{
     log::{WorkerEvent, WorkerEventType},
@@ -35,6 +35,32 @@ pub fn write_event(
     Ok(())
 }
 
+/// 从事件列表读取事件
+///
+/// # Errors
+/// 这个函数产生的错误来自Redis的错误[`redis::RedisError`]，以及读取内容不符合预期的错误
+///
+#[allow(clippy::cast_possible_wrap)]
+pub fn read_events(
+    redis_client: &Client,
+    start: usize,
+    limit: usize,
+) -> Result<Vec<WorkerEvent>, Error> {
+    redis_client
+        .get_connection()
+        .map_err(Error::Redis)?
+        .lrange(
+            worker_event_redis_key(),
+            start as isize,
+            (start + limit - 1) as isize,
+        )
+        .map_err(Error::Redis)?
+        .into_iter()
+        .map(|it| serde_json::from_str::<WorkerEvent>(&it))
+        .collect::<Result<Vec<_>, _>>()
+        .map_err(|e| Error::BadData(format!("failed to parse json to struct: ({e})")))
+}
+
 /// 使用`REDIS_CLIENT`尝试向事件列表写入新事件
 ///
 /// 很不卫生对吧，但是很好用

cuscuta-common/src/db/log/mod.rs

@@ -43,17 +43,18 @@ pub struct WorkerEvent {
 }
 
 /// Worker事件种类
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Clone, Copy)]
+#[repr(u8)]
 pub enum WorkerEventType {
     /// 最严重的错误（例如Worker离线）
-    Fatal,
+    Fatal = 9,
 
     /// Worker一般信息
-    Info,
+    Info = 1,
 
     /// Worker痕迹信息（例如拉取到任务，任务完成——会出现很多）
-    Trace,
+    Trace = 0,
 
     /// Worker警告（例如工作循环失败）
-    Warn,
+    Warn = 2,
 }

cuscuta-common/src/db/log/status.rs

@@ -1,7 +1,7 @@
 use std::borrow::Cow;
 
 use chrono::Utc;
-use redis::Client;
+use redis::{Client, TypedCommands};
 
 use crate::db::{
     job::{Job, SubQueue},
@@ -38,3 +38,28 @@ pub fn update_worker_status(
     pipe.exec(&mut connection).map_err(Error::Redis)?;
     Ok(())
 }
+
+/// 搜索Worker状态
+///
+/// # Errors
+/// 这个函数产生的错误来自Redis的错误[`redis::RedisError`]，以及读取内容不符合预期的错误
+///
+pub fn search_worker_status(
+    redis_client: &Client,
+) -> Result<Vec<(String, WorkerStatus<'_>)>, Error> {
+    let mut connection = redis_client.get_connection().map_err(Error::Redis)?;
+    connection
+        .scan_match::<_, String>(worker_status_redis_key("*"))
+        .map_err(Error::Redis)?
+        .collect::<Result<Vec<_>, _>>()
+        .map_err(Error::Redis)?
+        .into_iter()
+        .map(|it| connection.get(&it).map(|r| (it, r)))
+        .collect::<Result<Vec<_>, _>>()
+        .map_err(Error::Redis)?
+        .into_iter()
+        .filter_map(|(key, option_value)| option_value.map(|it| (key, it)))
+        .map(|(k, v)| serde_json::from_str::<WorkerStatus>(&v).map(|it| (k, it)))
+        .collect::<Result<Vec<_>, _>>()
+        .map_err(|e| Error::BadData(format!("failed to parse json to struct: ({e})")))
+}

cuscuta-worker/src/worker/pending_friend.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use cuscuta_common::{
     api::{
         self,
@@ -29,7 +31,7 @@ pub async fn try_add_friends(
     friends: &mut Vec<FriendInfo>,
 ) -> Result<(), Error> {
     let mut connection = redis_client.get_connection().map_err(Error::Redis)?;
-    let ids: Vec<_> = jobs
+    let mut ids: HashMap<_, _> = jobs
         .iter()
         .filter_map(|it| {
             let (JobState::Pending { friend_info, .. } | JobState::Finished { friend_info, .. }) =
@@ -44,10 +46,7 @@ pub async fn try_add_friends(
         let JobState::Pulled { start_timestamp } = job.state else {
             continue;
         };
-        let option_existing_ids = ids
-            .iter()
-            .find(|(code, _)| code == &job.essential.friend_code);
-        if let Some((_, existing_friend_info)) = option_existing_ids {
+        if let Some(existing_friend_info) = ids.get(&job.essential.friend_code) {
             let friend_info = existing_friend_info.clone();
             push_friend_info(&mut connection, &friend_info, job)?;
             job.state = JobState::Pending {
@@ -59,7 +58,9 @@ pub async fn try_add_friends(
             continue;
         }
         let friends_new =
-            match resolve_friend(config, bundle_data, account_row, user_id, token, job).await {
+            match try_modify_remote_friend(config, bundle_data, account_row, user_id, token, job)
+                .await
+            {
                 Ok(x) => x,
                 Err(e) => {
                     job.state = JobState::Failed {
@@ -94,6 +95,7 @@ pub async fn try_add_friends(
                 ));
             }
         };
+        ids.insert(job.essential.friend_code.clone(), friend_add.clone());
         *friends = friends_new;
         push_friend_info(&mut connection, &friend_add, job)?;
         job.essential.cursor_start = cursor.cast_signed() as i32;
@@ -106,7 +108,7 @@ pub async fn try_add_friends(
     Ok(())
 }
 
-async fn resolve_friend(
+async fn try_modify_remote_friend(
     config: &Config,
     bundle_data: &BundleData,
     account_row: &AccountRow,
@@ -139,7 +141,24 @@ async fn resolve_friend(
                     log::warn!(
                         "pending_friends: friend is already exist but cache is out-of-date! trying readd"
                     );
-                    delete_friend(config, bundle_data, account_row, user_id, token, job).await
+                    xxxxxx_safe_call_ex(
+                        config.worker_max_retry_count,
+                        config.worker_exponential_backoff_base_millis,
+                        config.worker_exponential_backoff_multiplier,
+                        config.worker_exponential_backoff_max_delay_millis,
+                        |it| it != StatusCode::TOO_MANY_REQUESTS,
+                        || {
+                            api::xxxxxx::api_delete_friend(
+                                bundle_data,
+                                &account_row.account_email,
+                                user_id,
+                                token,
+                                &job.essential.friend_code,
+                            )
+                        },
+                    )
+                    .await
+                    .map(|it| it.friends)
                 } else {
                     log::warn!("pending_friends: failed to add friend: {e}: code: {code}");
                     Err(e)
@@ -153,34 +172,6 @@ async fn resolve_friend(
     }
 }
 
-async fn delete_friend(
-    config: &Config,
-    bundle_data: &BundleData,
-    account_row: &AccountRow,
-    user_id: &str,
-    token: &str,
-    job: &Job,
-) -> Result<Vec<FriendInfo>, api::Error> {
-    xxxxxx_safe_call_ex(
-        config.worker_max_retry_count,
-        config.worker_exponential_backoff_base_millis,
-        config.worker_exponential_backoff_multiplier,
-        config.worker_exponential_backoff_max_delay_millis,
-        |it| it != StatusCode::TOO_MANY_REQUESTS,
-        || {
-            api::xxxxxx::api_delete_friend(
-                bundle_data,
-                &account_row.account_email,
-                user_id,
-                token,
-                &job.essential.friend_code,
-            )
-        },
-    )
-    .await
-    .map(|it| it.friends)
-}
-
 fn push_friend_info(
     connection: &mut Connection,
     friend_info: &FriendInfo,

cuscutactl/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "cuscutactl"
-version = "0.1.0"
+version = "0.1.1"
 edition = "2024"
 
 [dependencies]

cuscutactl/src/command.rs

@@ -1,4 +1,6 @@
-use clap::{Parser, Subcommand};
+use clap::{Parser, Subcommand, ValueEnum};
+use cuscuta_common::db::log::WorkerEventType;
+use serde::{Deserialize, Serialize};
 
 use crate::config::Config;
 
@@ -31,6 +33,13 @@ pub enum SubCommands {
         #[command(subcommand)]
         command: SubCommandAccounts,
     },
+
+    /// Worker stat query
+    #[command(visible_alias = "stat")]
+    Stats {
+        #[command(subcommand)]
+        command: SubCommandStats,
+    },
 }
 
 #[derive(Debug, Subcommand)]
@@ -157,3 +166,40 @@ pub enum SubCommandAccountsRate {
     /// Query current rating
     Query,
 }
+
+#[derive(Debug, Subcommand)]
+pub enum SubCommandStats {
+    /// Query worker state
+    Worker {},
+
+    /// Query worker event
+    Event {
+        #[arg(short, long, value_enum, default_value_t = ShowLevel::Info)]
+        show_level: ShowLevel,
+
+        #[arg(long, short)]
+        limit: usize,
+    },
+}
+
+#[derive(
+    Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum, Default, Serialize, Deserialize,
+)]
+pub enum ShowLevel {
+    Trace,
+    #[default]
+    Info,
+    Warn,
+    Fatal,
+}
+
+impl From<ShowLevel> for WorkerEventType {
+    fn from(value: ShowLevel) -> Self {
+        match value {
+            ShowLevel::Trace => Self::Trace,
+            ShowLevel::Info => Self::Info,
+            ShowLevel::Warn => Self::Warn,
+            ShowLevel::Fatal => Self::Fatal,
+        }
+    }
+}

cuscutactl/src/main.rs

@@ -54,6 +54,7 @@ mod config;
 mod doctor;
 mod jobs;
 mod kube;
+mod stats;
 mod url;
 
 trait Handler {
@@ -140,6 +141,7 @@ async fn main() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[allow(clippy::too_many_lines)]
 async fn run_command(
     cli: &Cli,
     pg_url: Result<String, String>,
@@ -228,6 +230,22 @@ async fn run_command(
                 }
             }
         }
+        SubCommands::Stats { command } => {
+            let redis_url = match redis_url {
+                Ok(x) => x,
+                Err(e) => {
+                    eprintln!("Configuration error: {e:?}");
+                    eprintln!("Hint: use --mode legacy --redis-url <URL>");
+                    bail!(e);
+                }
+            };
+            match command {
+                command::SubCommandStats::Worker {} => stats::worker(&redis_url)?,
+                command::SubCommandStats::Event { show_level, limit } => {
+                    stats::event(&redis_url, *show_level, *limit)?;
+                }
+            }
+        }
     }
     Ok(())
 }