#![allow(unused)]
use crate::cacheable::{CacheAble, CacheService};
use crate::errors::{CacheError, Error, Result};
use crate::utils::lock::DistributedLockManager;
use log::warn;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ErrorCategory {
Download,
Parse,
Auth,
RateLimit,
Other,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ErrorSeverity {
Minor,
Major,
Fatal,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ErrorRecord {
pub timestamp: u64,
pub category: ErrorCategory,
pub severity: ErrorSeverity,
pub message: String,
pub retry_count: usize,
}
impl CacheAble for ErrorRecord {
fn field() -> impl AsRef<str> {
"error_record".to_string()
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ErrorStats {
pub total_errors: usize,
pub success_count: usize,
pub consecutive_errors: usize,
pub errors_by_category: std::collections::HashMap<ErrorCategory, usize>,
pub last_error_time: Option<u64>,
pub last_success_time: Option<u64>,
pub is_task_terminated: bool,
pub is_module_terminated: bool,
}
impl CacheAble for ErrorStats {
fn field() -> impl AsRef<str> {
"error_stats".to_string()
}
}
impl ErrorStats {
pub fn error_rate(&self) -> f64 {
let total = self.total_errors + self.success_count;
if total == 0 {
0.0
} else {
self.total_errors as f64 / total as f64
}
}
pub fn health_score(&self) -> f64 {
let error_rate = self.error_rate();
let consecutive_penalty = (self.consecutive_errors as f64 * 0.1).min(0.5);
(1.0 - error_rate - consecutive_penalty).max(0.0)
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ModuleLocker {
pub is_locker: bool,
pub ts: u64,
}
impl CacheAble for ModuleLocker {
fn field() -> impl AsRef<str> {
"module_locker".to_string()
}
}
#[derive(Debug, Clone)]
pub enum ErrorDecision {
Continue,
RetryAfter(Duration),
Skip,
Terminate(String),
}
#[derive(Debug, Clone)]
pub struct ErrorTrackerConfig {
pub task_max_errors: usize,
pub module_max_errors: usize,
pub request_max_retries: usize,
pub parse_max_retries: usize,
pub enable_success_decay: bool,
pub success_decay_amount: usize,
pub enable_time_window: bool,
pub time_window_seconds: u64,
pub consecutive_error_threshold: usize,
pub error_ttl: u64,
}
impl Default for ErrorTrackerConfig {
fn default() -> Self {
Self {
task_max_errors: 100,
module_max_errors: 10,
request_max_retries: 3,
parse_max_retries: 3,
enable_success_decay: true,
success_decay_amount: 1,
enable_time_window: false,
time_window_seconds: 3600,
consecutive_error_threshold: 3,
error_ttl: 3600,
}
}
}
use dashmap::DashMap;
#[derive(Clone)]
pub struct StatusTracker {
cache_service: Arc<CacheService>,
config: ErrorTrackerConfig,
locker: Arc<DistributedLockManager>,
cache: Arc<DashMap<String, (ErrorStats, Instant)>>,
}
impl StatusTracker {
pub fn new(
cache_service: Arc<CacheService>,
config: ErrorTrackerConfig,
locker: Arc<DistributedLockManager>,
) -> Self {
Self {
cache_service,
config,
locker,
cache: Arc::new(DashMap::new()),
}
}
pub async fn record_download_error(
&self,
task_id: &str,
module_id: &str,
request_id: &str,
error: &Error,
) -> Result<ErrorDecision> {
let category = ErrorCategory::Download;
let severity = self.classify_error_severity(error);
let request_key = format!("request:{}:download", request_id);
let module_key = format!("module:{}:download", module_id);
let task_key = format!("task:{}:total", task_id);
let (request_res, module_res, task_res) = tokio::join!(
self.increment_error(&request_key, category, severity),
self.increment_error(&module_key, category, severity),
self.increment_error(&task_key, category, severity)
);
let request_errors = request_res?;
let module_errors = module_res?;
let task_errors = task_res?;
if severity == ErrorSeverity::Fatal {
return Ok(ErrorDecision::Terminate(format!(
"Fatal error encountered: {}",
error
)));
}
if task_errors >= self.config.task_max_errors {
return Ok(ErrorDecision::Terminate(format!(
"Task {} reached max errors: {}/{}",
task_id, task_errors, self.config.task_max_errors
)));
}
if module_errors >= self.config.module_max_errors {
self.release_module_locker(module_id).await;
return Ok(ErrorDecision::Terminate(format!(
"Module {} reached max errors: {}/{}",
module_id, module_errors, self.config.module_max_errors
)));
}
if request_errors >= self.config.request_max_retries {
return Ok(ErrorDecision::Skip);
}
let delay = match category {
ErrorCategory::RateLimit => Duration::from_secs(60),
ErrorCategory::Auth => Duration::from_secs(10),
_ => Duration::from_secs(2u64.pow(request_errors as u32).min(30)),
};
Ok(ErrorDecision::RetryAfter(delay))
}
pub async fn record_download_success(&self, request_id: &str) -> Result<()> {
let request_key = format!("request:{}:download", request_id);
if self.cache.contains_key(&request_key) {
self.cache.remove(&request_key);
}
Ok(())
}
pub async fn record_parse_error(
&self,
task_id: &str,
module_id: &str,
request_id: &str,
error: &Error,
) -> Result<ErrorDecision> {
let category = ErrorCategory::Parse;
let severity = self.classify_error_severity(error);
log::debug!(
"[ErrorTracker] record_parse_error: task_id={} module_id={} request_id={} error={}",
task_id,
module_id,
request_id,
error
);
let request_key = format!("request:{}:parse", request_id);
let module_key = format!("module:{}:parse", module_id);
let task_key = format!("task:{}:total", task_id);
let (request_res, module_res, task_res) = tokio::join!(
self.increment_error(&request_key, category, severity),
self.increment_error(&module_key, category, severity),
self.increment_error(&task_key, category, severity)
);
let request_errors = request_res?;
let module_errors = module_res?;
let task_errors = task_res?;
log::debug!(
"[ErrorTracker] parse error counts: task={}/{} module={}/{} request={}/{}",
task_errors,
self.config.task_max_errors,
module_errors,
self.config.module_max_errors,
request_errors,
self.config.parse_max_retries
);
if task_errors >= self.config.task_max_errors {
return Ok(ErrorDecision::Terminate(format!(
"Task {} reached max errors: {}/{}",
task_id, task_errors, self.config.task_max_errors
)));
}
if module_errors >= self.config.module_max_errors {
self.release_module_locker(module_id).await;
warn!(
"[ErrorTracker] module parse errors exceeded threshold: module_id={} errors={}/{}",
module_id, module_errors, self.config.module_max_errors
);
return Ok(ErrorDecision::Terminate(format!(
"Module {} reached max errors: {}/{}",
module_id, module_errors, self.config.module_max_errors
)));
}
if request_errors >= self.config.parse_max_retries {
return Ok(ErrorDecision::Skip);
}
Ok(ErrorDecision::RetryAfter(Duration::from_secs(1)))
}
pub async fn record_parse_success(&self, request_id: &str) -> Result<()> {
let request_key = format!("request:{}:parse", request_id);
if self.cache.contains_key(&request_key) {
self.cache.remove(&request_key);
}
Ok(())
}
pub async fn should_module_continue(&self, module_id: &str) -> Result<ErrorDecision> {
if let Some(entry) = self.cache.get(module_id) {
let (stats, _) = entry.value();
if stats.is_module_terminated {
return Ok(ErrorDecision::Terminate(format!(
"Module {} has been terminated (cached)",
module_id
)));
}
}
let download_key = format!("module:{}:download", module_id);
let parse_key = format!("module:{}:parse", module_id);
let (terminated_res, download_res, parse_res) = tokio::join!(
self.is_module_terminated(module_id),
self.get_error_count(&download_key),
self.get_error_count(&parse_key)
);
if terminated_res? {
warn!(
"[ErrorTracker] module already terminated: module_id={}",
module_id
);
return Ok(ErrorDecision::Terminate(format!(
"Module {} has been terminated",
module_id
)));
}
let download_errors = download_res?;
let parse_errors = parse_res?;
let total_errors = download_errors + parse_errors;
log::debug!(
"[ErrorTracker] should_module_continue: module_id={} total_errors={}/{} (download={}, parse={})",
module_id,
total_errors,
self.config.module_max_errors,
download_errors,
parse_errors
);
if total_errors >= self.config.module_max_errors {
self.mark_module_terminated(module_id).await?;
self.release_module_locker(module_id).await;
log::error!(
"[ErrorTracker] module exceeded error threshold: module_id={} total_errors={}/{} (download={}, parse={})",
module_id,
total_errors,
self.config.module_max_errors,
download_errors,
parse_errors
);
Ok(ErrorDecision::Terminate(format!(
"Module {} total errors: {}/{} (download: {}, parse: {})",
module_id,
total_errors,
self.config.module_max_errors,
download_errors,
parse_errors
)))
} else {
Ok(ErrorDecision::Continue)
}
}
pub async fn should_task_continue(&self, task_id: &str) -> Result<ErrorDecision> {
if let Some(entry) = self.cache.get(task_id) {
let (stats, _) = entry.value();
if stats.is_task_terminated {
return Ok(ErrorDecision::Terminate(format!(
"Task {} has been terminated (cached)",
task_id
)));
}
}
let task_key = format!("task:{}:total", task_id);
let (terminated_res, count_res) = tokio::join!(
self.is_task_terminated(task_id),
self.get_error_count(&task_key)
);
if terminated_res? {
return Ok(ErrorDecision::Terminate(format!(
"Task {} has been terminated",
task_id
)));
}
let task_errors = count_res?;
if task_errors >= self.config.task_max_errors {
self.mark_task_terminated(task_id).await?;
Ok(ErrorDecision::Terminate(format!(
"Task {} reached max errors: {}/{}",
task_id, task_errors, self.config.task_max_errors
)))
} else {
Ok(ErrorDecision::Continue)
}
}
pub async fn get_stats(&self, key: &str) -> Result<ErrorStats> {
if let Some(entry) = self.cache.get(key) {
let (stats, expires_at) = entry.value();
if Instant::now() < *expires_at {
return Ok(stats.clone());
} else {
}
} else {
}
self.get_stats_no_cache(key).await
}
pub async fn get_stats_no_cache(&self, key: &str) -> Result<ErrorStats> {
let stats_res = ErrorStats::sync(key, &self.cache_service).await;
match stats_res {
Ok(Some(stats)) => {
self.cache.insert(
key.to_string(),
(stats.clone(), Instant::now() + Duration::from_secs(10)),
);
Ok(stats)
}
Ok(None) => {
let stats = ErrorStats::default();
self.cache.insert(
key.to_string(),
(stats.clone(), Instant::now() + Duration::from_secs(5)),
);
Ok(stats)
}
Err(e) => Err(Error::from(e)),
}
}
pub async fn mark_task_terminated(&self, task_id: &str) -> Result<()> {
let cache_id = ErrorStats::cache_id(task_id, &self.cache_service);
self.locker
.with_lock(&cache_id, 10, Duration::from_secs(10), async {
let mut error_stats = self.get_stats_no_cache(task_id).await.unwrap_or_default();
error_stats.is_task_terminated = true;
error_stats.send(task_id, &self.cache_service).await
})
.await
.ok();
Ok(())
}
pub async fn is_task_terminated(&self, task_id: &str) -> Result<bool> {
match self.get_stats(task_id).await {
Ok(stats) => Ok(stats.is_task_terminated),
Err(e) => {
if let Some(cache_err) = e
.inner
.source
.as_ref()
.and_then(|s| s.downcast_ref::<CacheError>())
&& matches!(cache_err, CacheError::NotFound)
{
return Ok(false);
}
warn!(
"[ErrorTracker] task terminated check failed: task_id={}, error: {}",
task_id, e
);
Err(e)
}
}
}
pub async fn mark_module_terminated(&self, module_id: &str) -> Result<()> {
let cache_id = ErrorStats::cache_id(module_id, &self.cache_service);
self.locker
.with_lock(&cache_id, 10, Duration::from_secs(10), async {
let mut error_stats = self.get_stats_no_cache(module_id).await.unwrap_or_default();
error_stats.is_module_terminated = true;
error_stats.send(module_id, &self.cache_service).await
})
.await
.ok();
Ok(())
}
pub async fn is_module_terminated(&self, module_id: &str) -> Result<bool> {
match self.get_stats(module_id).await {
Ok(stats) => Ok(stats.is_module_terminated),
Err(e) => {
if let Some(cache_err) = e
.inner
.source
.as_ref()
.and_then(|s| s.downcast_ref::<CacheError>())
&& matches!(cache_err, CacheError::NotFound)
{
return Ok(false);
}
warn!(
"[ErrorTracker] module terminated check failed: module_id={}, error: {}",
module_id, e
);
Err(e)
}
}
}
async fn increment_error(
&self,
key: &str,
category: ErrorCategory,
_severity: ErrorSeverity,
) -> Result<usize> {
if key.starts_with("request:") {
let count_key = format!("{}:total_errors", key);
let total = self.cache_service.incr(&count_key, 1).await? as usize;
return Ok(total);
}
let count_key = format!("{}:total_errors", key);
let total = self.cache_service.incr(&count_key, 1).await? as usize;
let consecutive_key = format!("{}:consecutive_errors", key);
let _ = self.cache_service.incr(&consecutive_key, 1).await;
Ok(total)
}
async fn decrement_error(&self, key: &str, amount: usize) -> Result<()> {
let count_key = format!("{}:total_errors", key);
let _ = self.cache_service.incr(&count_key, -(amount as i64)).await;
let consecutive_key = format!("{}:consecutive_errors", key);
let _ = self.cache_service.set(&consecutive_key, b"0", None).await;
Ok(())
}
async fn increment_success(&self, key: &str) -> Result<()> {
let mut stats = self.get_stats(key).await.unwrap_or_default();
stats.success_count += 1;
stats.last_success_time = Some(Self::now());
self.save_stats_with_ttl(key, &stats).await.ok();
Ok(())
}
async fn reset_consecutive_errors(&self, key: &str) -> Result<()> {
let mut stats = self.get_stats(key).await.unwrap_or_default();
stats.consecutive_errors = 0;
self.save_stats_with_ttl(key, &stats).await.ok();
Ok(())
}
async fn get_error_count(&self, key: &str) -> Result<usize> {
let count_key = format!("{}:total_errors", key);
if let Some(val) = self.cache_service.get(&count_key).await? {
let s = String::from_utf8(val).unwrap_or_default();
let count = s.parse::<usize>().unwrap_or(0);
return Ok(count);
}
Ok(0)
}
async fn save_stats(&self, key: &str, stats: &ErrorStats) -> Result<()> {
self.save_stats_with_ttl(key, stats).await
}
async fn save_stats_with_ttl(&self, key: &str, stats: &ErrorStats) -> Result<()> {
self.cache.insert(
key.to_string(),
(stats.clone(), Instant::now() + Duration::from_secs(10)),
);
stats
.send(key, &self.cache_service)
.await
.map_err(Error::from)
}
fn classify_error_severity(&self, error: &Error) -> ErrorSeverity {
let error_str = error.to_string().to_lowercase();
if error_str.contains("auth") || error_str.contains("unauthorized") {
ErrorSeverity::Fatal
} else if error_str.contains("rate limit") || error_str.contains("too many") {
ErrorSeverity::Major
} else {
ErrorSeverity::Minor
}
}
fn now() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs()
}
pub async fn lock_module(&self, module_id: &str) {
let ts = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap()
.as_secs();
let module_locker = ModuleLocker {
is_locker: true,
ts,
};
module_locker
.send(module_id, &self.cache_service)
.await
.ok();
}
pub async fn is_module_locker(&self, module_id: &str, ttl: u64) -> bool {
ModuleLocker::sync(module_id, &self.cache_service)
.await
.ok()
.flatten()
.is_some_and(|locker| {
if locker.is_locker {
let now = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap()
.as_secs();
now - locker.ts <= ttl
} else {
false
}
})
}
pub async fn release_module_locker(&self, module_id: &str) {
ModuleLocker::delete(module_id, &self.cache_service)
.await
.ok();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_error_stats_calculations() {
let mut stats = ErrorStats::default();
stats.total_errors = 3;
stats.success_count = 7;
assert_eq!(stats.error_rate(), 0.3);
assert!(stats.health_score() > 0.6);
stats.consecutive_errors = 5;
assert!(stats.health_score() < 0.5);
}
}