From e81211c88e2be5a41970bea6b4c487d2f176ee45 Mon Sep 17 00:00:00 2001 From: Valmo Date: Tue, 30 Jun 2026 18:07:47 -0300 Subject: [PATCH] Refactored redis component to use a centralized postgres database with coordinator instead of descentralized sqlite shitty databases --- Cargo.lock | 30 ++ Cargo.toml | 1 + Dockerfile | 2 +- docker-compose.redis-rtsp.yml | 26 +- docker-compose.redis-rtsp_pro_max.yml | 34 +- docker-compose.redis-workers.yml | 24 +- src/config.rs | 12 - src/db.rs | 79 +++ src/redis_rtsp.rs | 674 ++++++++++++++++++++------ src/terminal.rs | 165 +++++++ 10 files changed, 854 insertions(+), 193 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3cdbb26..c07327f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -64,6 +64,7 @@ version = "0.1.0" dependencies = [ "anyhow", "dotenvy", + "terminal_size", "tokio", "tokio-postgres", ] @@ -255,6 +256,12 @@ dependencies = [ "libc", ] +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "lock_api" version = "0.4.14" @@ -454,6 +461,19 @@ dependencies = [ "bitflags", ] +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys", +] + [[package]] name = "rustversion" version = "1.0.22" @@ -566,6 +586,16 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "terminal_size" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "230a1b821ccbd75b185820a1f1ff7b14d21da1e442e22c0863ea5f08771a8874" +dependencies = [ + "rustix", + "windows-sys", +] + [[package]] name = "tinyvec" version = "1.11.0" diff --git a/Cargo.toml b/Cargo.toml index 6ddad19..17885fa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,3 +8,4 @@ anyhow = "1" dotenvy = "0.15" tokio = { version = "1", features = ["io-util", "macros", "net", "process", "rt-multi-thread", "signal", "time"] } tokio-postgres = "0.7" +terminal_size = "0.4" diff --git a/Dockerfile b/Dockerfile index cb0543c..3b732dd 100644 --- a/Dockerfile +++ b/Dockerfile @@ -12,7 +12,7 @@ RUN cargo build --release FROM debian:bookworm-slim RUN apt-get update \ - && apt-get install -y --no-install-recommends ca-certificates ffmpeg redis-tools sqlite3 \ + && apt-get install -y --no-install-recommends ca-certificates ffmpeg redis-tools \ && rm -rf /var/lib/apt/lists/* WORKDIR /app diff --git a/docker-compose.redis-rtsp.yml b/docker-compose.redis-rtsp.yml index bdeca60..658c123 100644 --- a/docker-compose.redis-rtsp.yml +++ b/docker-compose.redis-rtsp.yml @@ -2,21 +2,38 @@ services: redis: image: redis:7-alpine command: ["redis-server", "--appendonly", "yes", "--bind", "0.0.0.0", "--protected-mode", "no"] - ports: - - "6379:6379" volumes: - camfinder-redis:/data + postgres: + image: postgres:16-alpine + environment: + POSTGRES_DB: camfinder + POSTGRES_USER: camfinder + POSTGRES_PASSWORD: camfinder + volumes: + - camfinder-postgres:/var/lib/postgresql/data + healthcheck: + test: ["CMD-SHELL", "pg_isready -U camfinder -d camfinder"] + interval: 5s + timeout: 5s + retries: 10 + start_period: 5s + coordinator: build: . command: ["redis-coordinator"] environment: REDIS_URL: redis://redis:6379/0 + DATABASE_URL: postgres://camfinder:camfinder@postgres:5432/camfinder RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_INPUT_CSV: /app/camfinder-open-rtsp.csv RTSP_WORKER_COUNT: 2 depends_on: - - redis + redis: + condition: service_started + postgres: + condition: service_healthy volumes: - ./:/app @@ -27,7 +44,6 @@ services: REDIS_URL: redis://redis:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 depends_on: - redis @@ -42,7 +58,6 @@ services: REDIS_URL: redis://redis:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 depends_on: - redis @@ -52,3 +67,4 @@ services: volumes: camfinder-redis: + camfinder-postgres: diff --git a/docker-compose.redis-rtsp_pro_max.yml b/docker-compose.redis-rtsp_pro_max.yml index d3482f6..589b633 100644 --- a/docker-compose.redis-rtsp_pro_max.yml +++ b/docker-compose.redis-rtsp_pro_max.yml @@ -2,21 +2,38 @@ services: redis: image: redis:7-alpine command: ["redis-server", "--appendonly", "yes", "--bind", "0.0.0.0", "--protected-mode", "no"] - ports: - - "6379:6379" volumes: - camfinder-redis:/data + postgres: + image: postgres:16-alpine + environment: + POSTGRES_DB: camfinder + POSTGRES_USER: camfinder + POSTGRES_PASSWORD: camfinder + volumes: + - camfinder-postgres:/var/lib/postgresql/data + healthcheck: + test: ["CMD-SHELL", "pg_isready -U camfinder -d camfinder"] + interval: 5s + timeout: 5s + retries: 10 + start_period: 5s + coordinator: build: . command: ["redis-coordinator"] environment: REDIS_URL: redis://redis:6379/0 + DATABASE_URL: postgres://camfinder:camfinder@postgres:5432/camfinder RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_INPUT_CSV: /app/camfinder-open-rtsp.csv - RTSP_WORKER_COUNT: 2 + RTSP_WORKER_COUNT: 8 depends_on: - - redis + redis: + condition: service_started + postgres: + condition: service_healthy volumes: - ./:/app:ro @@ -27,7 +44,6 @@ services: REDIS_URL: redis://redis:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 depends_on: - redis @@ -42,7 +58,6 @@ services: REDIS_URL: redis://redis:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 depends_on: - worker-1 @@ -56,7 +71,6 @@ services: REDIS_URL: redis://redis:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 depends_on: - worker-2 @@ -70,7 +84,6 @@ services: REDIS_URL: redis://redis:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 depends_on: - worker-3 @@ -84,7 +97,6 @@ services: REDIS_URL: redis://redis:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 depends_on: - worker-4 @@ -98,7 +110,6 @@ services: REDIS_URL: redis://redis:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 depends_on: - worker-5 @@ -112,7 +123,6 @@ services: REDIS_URL: redis://redis:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 depends_on: - worker-6 @@ -126,7 +136,6 @@ services: REDIS_URL: redis://redis:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 depends_on: - worker-7 @@ -134,3 +143,4 @@ services: - ./dump:/app/dump volumes: camfinder-redis: + camfinder-postgres: diff --git a/docker-compose.redis-workers.yml b/docker-compose.redis-workers.yml index 7408b33..2bdf956 100644 --- a/docker-compose.redis-workers.yml +++ b/docker-compose.redis-workers.yml @@ -3,10 +3,9 @@ services: build: . command: ["redis-worker"] environment: - REDIS_URL: redis://192.168.15.101:6379/0 + REDIS_URL: redis://192.168.15.100:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 restart: unless-stopped volumes: @@ -16,10 +15,9 @@ services: build: . command: ["redis-worker"] environment: - REDIS_URL: redis://192.168.15.101:6379/0 + REDIS_URL: redis://192.168.15.100:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 restart: unless-stopped volumes: @@ -29,10 +27,9 @@ services: build: . command: ["redis-worker"] environment: - REDIS_URL: redis://192.168.15.101:6379/0 + REDIS_URL: redis://192.168.15.100:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 restart: unless-stopped volumes: @@ -42,10 +39,9 @@ services: build: . command: ["redis-worker"] environment: - REDIS_URL: redis://192.168.15.101:6379/0 + REDIS_URL: redis://192.168.15.100:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 restart: unless-stopped volumes: @@ -55,10 +51,9 @@ services: build: . command: ["redis-worker"] environment: - REDIS_URL: redis://192.168.15.101:6379/0 + REDIS_URL: redis://192.168.15.100:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 restart: unless-stopped volumes: @@ -68,10 +63,9 @@ services: build: . command: ["redis-worker"] environment: - REDIS_URL: redis://192.168.15.101:6379/0 + REDIS_URL: redis://192.168.15.100:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 restart: unless-stopped volumes: @@ -81,10 +75,9 @@ services: build: . command: ["redis-worker"] environment: - REDIS_URL: redis://192.168.15.101:6379/0 + REDIS_URL: redis://192.168.15.100:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 restart: unless-stopped volumes: @@ -94,10 +87,9 @@ services: build: . command: ["redis-worker"] environment: - REDIS_URL: redis://192.168.15.101:6379/0 + REDIS_URL: redis://192.168.15.100:6379/0 RTSP_QUEUE_NAME: camfinder:rtsp:queue RTSP_SNAPSHOT_DIR: /app/dump/feeds - RTSP_SQLITE_PATH: /app/dump/rtsp-results.sqlite RTSP_CAPTURE_TIMEOUT_MS: 15000 restart: unless-stopped volumes: diff --git a/src/config.rs b/src/config.rs index af113dc..2622da7 100644 --- a/src/config.rs +++ b/src/config.rs @@ -27,7 +27,6 @@ pub struct Config { pub redis_worker_count: usize, pub redis_queue_block_size: usize, pub snapshot_dir: PathBuf, - pub sqlite_path: PathBuf, pub capture_timeout_ms: u64, pub port: u16, pub timeout_ms: u64, @@ -73,9 +72,6 @@ impl Config { snapshot_dir: cli .snapshot_dir .unwrap_or_else(|| PathBuf::from(read_env("RTSP_SNAPSHOT_DIR", "dump/feeds"))), - sqlite_path: cli.sqlite_path.unwrap_or_else(|| { - PathBuf::from(read_env("RTSP_SQLITE_PATH", "dump/rtsp-results.sqlite")) - }), capture_timeout_ms: cli .capture_timeout_ms .or_else(|| parse_u64("RTSP_CAPTURE_TIMEOUT_MS", 15000).ok()) @@ -123,7 +119,6 @@ struct CliArgs { redis_worker_count: Option, redis_queue_block_size: Option, snapshot_dir: Option, - sqlite_path: Option, capture_timeout_ms: Option, database_url: Option, coordinator_bind: Option, @@ -197,13 +192,6 @@ impl CliArgs { "--snapshot-dir", )?)); } - "--sqlite-path" => { - cli.sqlite_path = Some(PathBuf::from(next_value( - args, - &mut index, - "--sqlite-path", - )?)); - } "--capture-timeout-ms" => { cli.capture_timeout_ms = Some(parse_u64_literal(&next_value( args, diff --git a/src/db.rs b/src/db.rs index 6d5fd50..c49547f 100644 --- a/src/db.rs +++ b/src/db.rs @@ -69,6 +69,29 @@ impl Database { CREATE INDEX IF NOT EXISTS rtsp_probe_results_status_idx ON rtsp_probe_results (status); + + CREATE TABLE IF NOT EXISTS rtsp_fetch_results ( + id BIGSERIAL PRIMARY KEY, + checked_at TIMESTAMPTZ NOT NULL, + checked_at_ms BIGINT NOT NULL, + target_ip TEXT NOT NULL, + target_port INTEGER NOT NULL, + rtsp_path TEXT NOT NULL, + rtsp_url TEXT NOT NULL, + screenshot_path TEXT, + status TEXT NOT NULL, + error TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + ); + + CREATE INDEX IF NOT EXISTS rtsp_fetch_results_checked_at_idx + ON rtsp_fetch_results (checked_at); + + CREATE INDEX IF NOT EXISTS rtsp_fetch_results_target_ip_idx + ON rtsp_fetch_results (target_ip); + + CREATE INDEX IF NOT EXISTS rtsp_fetch_results_status_idx + ON rtsp_fetch_results (status); "#, ) .await @@ -252,6 +275,62 @@ impl Database { Ok(()) } + pub async fn insert_rtsp_fetch_result( + &self, + checked_at_ms: i64, + target_ip: &str, + target_port: u16, + rtsp_path: &str, + rtsp_url: &str, + screenshot_path: Option<&str>, + status: &str, + error: Option<&str>, + ) -> Result<()> { + let target_port = i32::from(target_port); + + self.client + .execute( + r#" + INSERT INTO rtsp_fetch_results ( + checked_at, + checked_at_ms, + target_ip, + target_port, + rtsp_path, + rtsp_url, + screenshot_path, + status, + error + ) + VALUES ( + TO_TIMESTAMP($1::DOUBLE PRECISION / 1000.0), + $1, + $2, + $3, + $4, + $5, + $6, + $7, + $8 + ) + "#, + &[ + &checked_at_ms, + &target_ip, + &target_port, + &rtsp_path, + &rtsp_url, + &screenshot_path, + &status, + &error, + ], + ) + .await + .context("failed to insert RTSP fetch result")?; + + Ok(()) + } + pub async fn summarize_chunks(&self) -> Result<(usize, usize)> { let row = self .client diff --git a/src/redis_rtsp.rs b/src/redis_rtsp.rs index 304147a..c058496 100644 --- a/src/redis_rtsp.rs +++ b/src/redis_rtsp.rs @@ -1,6 +1,7 @@ use crate::{ config::Config, - terminal::{ScreenState, TerminalUi}, + db::Database, + terminal::{RedisWorkerScreenState, ScreenState, TerminalUi}, }; use anyhow::{anyhow, Context, Result}; use std::{ @@ -20,6 +21,9 @@ use tokio::{ const DEFAULT_DONE_SUFFIX: &str = ":done"; const DEFAULT_SENTINEL: &str = "__CAMFINDER_DONE__"; +const GLOBAL_PROGRESS_REFRESH: Duration = Duration::from_secs(5); +const UI_REFRESH: Duration = Duration::from_millis(250); +const RESULT_QUEUE_SUFFIX: &str = ":results"; pub async fn run_redis_coordinator( config: Config, @@ -34,8 +38,22 @@ pub async fn run_redis_coordinator( )); } + let database = Database::open(&config.database_url).await?; let done_key = format!("{}{}", config.redis_queue_name, DEFAULT_DONE_SUFFIX); - clear_queue(&config.redis_url, &config.redis_queue_name, &done_key).await?; + let result_queue = result_queue_name(&config.redis_queue_name); + let total_key = total_ips_key(&config.redis_queue_name); + let completed_key = completed_ips_key(&config.redis_queue_name); + clear_queue( + &config.redis_url, + &config.redis_queue_name, + &result_queue, + &done_key, + &total_key, + &completed_key, + ) + .await?; + set_redis_key(&config.redis_url, &total_key, ips.len()).await?; + set_redis_key(&config.redis_url, &completed_key, 0).await?; enqueue_ips( &config.redis_url, &config.redis_queue_name, @@ -63,30 +81,74 @@ pub async fn run_redis_coordinator( ); } + consume_result_queue( + &config.redis_url, + &config.redis_queue_name, + &result_queue, + &completed_key, + ips.len(), + &database, + ) + .await?; + Ok(()) } pub async fn run_redis_worker( config: Config, - _ui: Arc, + ui: Arc, interrupted: Arc, ) -> Result<()> { - ensure_output_dirs(&config.snapshot_dir, &config.sqlite_path)?; - init_sqlite(&config.sqlite_path).await?; + ensure_snapshot_dir(&config.snapshot_dir)?; let done_key = format!("{}{}", config.redis_queue_name, DEFAULT_DONE_SUFFIX); + let result_queue = result_queue_name(&config.redis_queue_name); + let total_key = total_ips_key(&config.redis_queue_name); + let completed_key = completed_ips_key(&config.redis_queue_name); + ensure_total_ips_key(&config.redis_url, &total_key, &config.redis_input_csv).await?; + let started_at = Instant::now(); let mut processed = 0usize; let mut captured = 0usize; let paths = normalize_paths(&config.rtsp_paths); let credentials = load_rtsp_credentials(Path::new("/app/main_ids.json"))?; + let worker_name = format!("{}-{}", config.worker_name, std::process::id()); + let mut global_stats = read_global_progress( + &config.redis_url, + &config.redis_queue_name, + &total_key, + &completed_key, + ) + .await + .unwrap_or_default(); + let mut last_global_refresh = Instant::now(); + let mut last_ui_render = Instant::now() + .checked_sub(UI_REFRESH) + .unwrap_or_else(Instant::now); - println!( - "worker pronto para {} paths, {} credenciais e fila {}", - paths.len(), - credentials.len(), - config.redis_queue_name - ); + if ui.is_enabled() { + render_worker_progress( + &ui, + &worker_name, + &config.redis_queue_name, + "aguardando bloco", + "", + "iniciando", + processed, + captured, + &global_stats, + 0, + 0, + started_at, + ); + } else { + println!( + "worker pronto para {} paths, {} credenciais e fila {}", + paths.len(), + credentials.len(), + config.redis_queue_name + ); + } loop { if interrupted.load(Ordering::SeqCst) { @@ -96,7 +158,14 @@ pub async fn run_redis_worker( match pop_queue_item(&config.redis_url, &config.redis_queue_name).await? { Some(item) if item == DEFAULT_SENTINEL => break, Some(block) => { - for ip in parse_ip_block(&block) { + let block_ips = parse_ip_block(&block); + let block_total = block_ips + .len() + .saturating_mul(paths.len()) + .saturating_mul(credentials.len()); + let mut block_done = 0usize; + + for ip in block_ips { for rtsp_path in &paths { if interrupted.load(Ordering::SeqCst) { break; @@ -154,24 +223,69 @@ pub async fn run_redis_worker( ) }; - insert_sqlite_result( - &config.sqlite_path, - checked_at, - &ip, - &rtsp_path, - &attempt, - status, - screenshot_path.as_deref(), - error.as_deref(), + push_fetch_result( + &config.redis_url, + &result_queue, + &RtspFetchResult { + checked_at, + target_ip: ip.clone(), + target_port: config.port, + rtsp_path: rtsp_path.clone(), + rtsp_url: attempt.clone(), + screenshot_path: screenshot_path.clone(), + status, + error, + }, ) .await?; processed += 1; - println!( - "{} -> {}", - redact_rtsp_url(&attempt), - screenshot_path.unwrap_or_else(|| "sem imagem".to_string()) - ); + block_done = block_done.saturating_add(1); + + let output_label = + screenshot_path.unwrap_or_else(|| "sem imagem".to_string()); + if ui.is_enabled() { + let now = Instant::now(); + if now.duration_since(last_global_refresh) + >= GLOBAL_PROGRESS_REFRESH + { + if let Ok(stats) = read_global_progress( + &config.redis_url, + &config.redis_queue_name, + &total_key, + &completed_key, + ) + .await + { + global_stats = stats; + } + last_global_refresh = now; + } + + if now.duration_since(last_ui_render) >= UI_REFRESH { + render_worker_progress( + &ui, + &worker_name, + &config.redis_queue_name, + &ip, + rtsp_path, + &format!( + "{} -> {}", + redact_rtsp_url(&attempt), + output_label + ), + processed, + captured, + &global_stats, + block_done, + block_total, + started_at, + ); + last_ui_render = now; + } + } else { + println!("{} -> {}", redact_rtsp_url(&attempt), output_label); + } if attempt_captured { captured += 1; @@ -188,9 +302,45 @@ pub async fn run_redis_worker( continue; } } + + increment_redis_key(&config.redis_url, &completed_key).await?; + global_stats.completed_ips = global_stats.completed_ips.saturating_add(1); + } + + if ui.is_enabled() { + render_worker_progress( + &ui, + &worker_name, + &config.redis_queue_name, + "aguardando bloco", + "", + "bloco finalizado", + processed, + captured, + &global_stats, + block_total, + block_total, + started_at, + ); } } None => { + if ui.is_enabled() { + render_worker_progress( + &ui, + &worker_name, + &config.redis_queue_name, + "fila vazia", + "", + "aguardando done", + processed, + captured, + &global_stats, + 0, + 0, + started_at, + ); + } if redis_key_exists(&config.redis_url, &done_key).await? { break; } @@ -200,10 +350,28 @@ pub async fn run_redis_worker( } let elapsed = started_at.elapsed().as_secs_f64(); - println!( - "worker finalizado em {:.2}s, tentativas {}, capturas {}", - elapsed, processed, captured - ); + if ui.is_enabled() { + render_worker_progress( + &ui, + &worker_name, + &config.redis_queue_name, + "finalizado", + "", + &format!("tentativas {}, capturas {}", processed, captured), + processed, + captured, + &global_stats, + 1, + 1, + started_at, + ); + ui.finish(); + } else { + println!( + "worker finalizado em {:.2}s, tentativas {}, capturas {}", + elapsed, processed, captured + ); + } Ok(()) } @@ -365,6 +533,25 @@ struct RtspProbe { error: Option, } +#[derive(Debug, Clone, Default)] +struct GlobalProgress { + completed_ips: usize, + total_ips: usize, + queue_remaining_ips: Option, +} + +#[derive(Debug, Clone)] +struct RtspFetchResult { + checked_at: i64, + target_ip: String, + target_port: u16, + rtsp_path: String, + rtsp_url: String, + screenshot_path: Option, + status: String, + error: Option, +} + impl RtspProbe { fn is_connectable(&self) -> bool { matches!( @@ -423,8 +610,26 @@ async fn redis_key_exists(redis_url: &str, key: &str) -> Result { Ok(output.trim() == "1") } -async fn clear_queue(redis_url: &str, queue_name: &str, done_key: &str) -> Result<()> { - let _ = redis_cli(redis_url, &["DEL", queue_name, done_key]).await?; +async fn clear_queue( + redis_url: &str, + queue_name: &str, + result_queue: &str, + done_key: &str, + total_key: &str, + completed_key: &str, +) -> Result<()> { + let _ = redis_cli( + redis_url, + &[ + "DEL", + queue_name, + result_queue, + done_key, + total_key, + completed_key, + ], + ) + .await?; Ok(()) } @@ -433,6 +638,294 @@ async fn mark_done(redis_url: &str, done_key: &str) -> Result<()> { Ok(()) } +async fn set_redis_key(redis_url: &str, key: &str, value: usize) -> Result<()> { + let value = value.to_string(); + let _ = redis_cli(redis_url, &["SET", key, &value]).await?; + Ok(()) +} + +async fn set_redis_key_if_missing(redis_url: &str, key: &str, value: usize) -> Result<()> { + let value = value.to_string(); + let _ = redis_cli(redis_url, &["SETNX", key, &value]).await?; + Ok(()) +} + +async fn increment_redis_key(redis_url: &str, key: &str) -> Result<()> { + let _ = redis_cli(redis_url, &["INCR", key]).await?; + Ok(()) +} + +async fn ensure_total_ips_key(redis_url: &str, total_key: &str, input_csv: &str) -> Result<()> { + let total = load_ip_list(input_csv)?.len(); + set_redis_key_if_missing(redis_url, total_key, total).await +} + +async fn read_global_progress( + redis_url: &str, + queue_name: &str, + total_key: &str, + completed_key: &str, +) -> Result { + let script = r#" + local queue = KEYS[1] + local total_key = KEYS[2] + local completed_key = KEYS[3] + local total = tonumber(redis.call('GET', total_key) or '0') + local completed = tonumber(redis.call('GET', completed_key) or '0') + local n = redis.call('LLEN', queue) + local remaining = 0 + for i = 0, n - 1 do + local item = redis.call('LINDEX', queue, i) + if item ~= '__CAMFINDER_DONE__' then + for _ in string.gmatch(item, '[^\n]+') do + remaining = remaining + 1 + end + end + end + return {completed, total, remaining} + "#; + let output = redis_cli( + redis_url, + &[ + "--raw", + "EVAL", + script, + "3", + queue_name, + total_key, + completed_key, + ], + ) + .await?; + let values = output + .lines() + .filter_map(|line| line.trim().parse::().ok()) + .collect::>(); + Ok(GlobalProgress { + completed_ips: values.get(0).copied().unwrap_or_default(), + total_ips: values.get(1).copied().unwrap_or_default(), + queue_remaining_ips: values.get(2).copied(), + }) +} + +fn total_ips_key(queue_name: &str) -> String { + format!("{queue_name}:total_ips") +} + +fn completed_ips_key(queue_name: &str) -> String { + format!("{queue_name}:completed_ips") +} + +fn result_queue_name(queue_name: &str) -> String { + format!("{queue_name}{RESULT_QUEUE_SUFFIX}") +} + +async fn push_fetch_result( + redis_url: &str, + result_queue: &str, + result: &RtspFetchResult, +) -> Result<()> { + let serialized = serialize_fetch_result(result); + let _ = redis_cli(redis_url, &["RPUSH", result_queue, &serialized]).await?; + Ok(()) +} + +async fn pop_fetch_result(redis_url: &str, result_queue: &str) -> Result> { + let output = redis_cli(redis_url, &["--raw", "LPOP", result_queue]).await?; + let value = output.trim_end_matches('\n'); + if value.is_empty() { + return Ok(None); + } + deserialize_fetch_result(value).map(Some) +} + +async fn consume_result_queue( + redis_url: &str, + queue_name: &str, + result_queue: &str, + completed_key: &str, + total_ips: usize, + database: &Database, +) -> Result<()> { + let mut inserted = 0usize; + loop { + let mut drained_any = false; + while let Some(result) = pop_fetch_result(redis_url, result_queue).await? { + database + .insert_rtsp_fetch_result( + result.checked_at, + &result.target_ip, + result.target_port, + &result.rtsp_path, + &result.rtsp_url, + result.screenshot_path.as_deref(), + &result.status, + result.error.as_deref(), + ) + .await?; + inserted += 1; + drained_any = true; + } + + let completed = redis_usize(redis_url, completed_key) + .await? + .unwrap_or_default(); + let pending_results = redis_list_len(redis_url, result_queue).await?; + if completed >= total_ips && pending_results == 0 { + println!( + "writer finalizado: {} resultados gravados em Postgres para {}", + inserted, queue_name + ); + break; + } + + if !drained_any { + sleep(Duration::from_millis(500)).await; + } + } + + Ok(()) +} + +async fn redis_usize(redis_url: &str, key: &str) -> Result> { + let output = redis_cli(redis_url, &["--raw", "GET", key]).await?; + let value = output.trim(); + if value.is_empty() { + return Ok(None); + } + value + .parse::() + .map(Some) + .with_context(|| format!("invalid Redis integer for {key}: {value}")) +} + +async fn redis_list_len(redis_url: &str, key: &str) -> Result { + let output = redis_cli(redis_url, &["--raw", "LLEN", key]).await?; + output + .trim() + .parse::() + .with_context(|| format!("invalid Redis LLEN for {key}")) +} + +fn serialize_fetch_result(result: &RtspFetchResult) -> String { + [ + result.checked_at.to_string(), + result.target_port.to_string(), + encode_field(&result.target_ip), + encode_field(&result.rtsp_path), + encode_field(&result.rtsp_url), + encode_optional_field(result.screenshot_path.as_deref()), + encode_field(&result.status), + encode_optional_field(result.error.as_deref()), + ] + .join("\t") +} + +fn deserialize_fetch_result(value: &str) -> Result { + let parts = value.split('\t').collect::>(); + if parts.len() != 8 { + return Err(anyhow!( + "invalid RTSP fetch result field count: expected 8, got {}", + parts.len() + )); + } + Ok(RtspFetchResult { + checked_at: parts[0] + .parse::() + .with_context(|| format!("invalid checked_at: {}", parts[0]))?, + target_port: parts[1] + .parse::() + .with_context(|| format!("invalid target_port: {}", parts[1]))?, + target_ip: decode_field(parts[2])?, + rtsp_path: decode_field(parts[3])?, + rtsp_url: decode_field(parts[4])?, + screenshot_path: decode_optional_field(parts[5])?, + status: decode_field(parts[6])?, + error: decode_optional_field(parts[7])?, + }) +} + +fn encode_optional_field(value: Option<&str>) -> String { + value.map(encode_field).unwrap_or_default() +} + +fn decode_optional_field(value: &str) -> Result> { + if value.is_empty() { + Ok(None) + } else { + decode_field(value).map(Some) + } +} + +fn encode_field(value: &str) -> String { + let mut encoded = String::with_capacity(value.len()); + for byte in value.as_bytes() { + match byte { + b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'.' | b'-' | b'_' | b'/' | b':' | b'@' => { + encoded.push(*byte as char) + } + _ => { + let _ = write!(&mut encoded, "%{byte:02X}"); + } + } + } + encoded +} + +fn decode_field(value: &str) -> Result { + let bytes = value.as_bytes(); + let mut decoded = Vec::with_capacity(bytes.len()); + let mut index = 0usize; + while index < bytes.len() { + if bytes[index] == b'%' { + if index + 2 >= bytes.len() { + return Err(anyhow!("invalid percent escape in Redis result field")); + } + let hex = std::str::from_utf8(&bytes[index + 1..index + 3]) + .context("invalid percent escape UTF-8")?; + let value = u8::from_str_radix(hex, 16) + .with_context(|| format!("invalid percent escape: %{hex}"))?; + decoded.push(value); + index += 3; + } else { + decoded.push(bytes[index]); + index += 1; + } + } + String::from_utf8(decoded).context("decoded Redis result field is not valid UTF-8") +} + +fn render_worker_progress( + ui: &TerminalUi, + worker_name: &str, + queue_name: &str, + current_ip: &str, + current_path: &str, + current_status: &str, + processed: usize, + captured: usize, + global_stats: &GlobalProgress, + block_done: usize, + block_total: usize, + started_at: Instant, +) { + ui.render_redis_worker(&RedisWorkerScreenState { + worker_name: worker_name.to_string(), + queue_name: queue_name.to_string(), + current_ip: current_ip.to_string(), + current_path: current_path.to_string(), + current_status: current_status.to_string(), + processed, + captured, + global_done: global_stats.completed_ips, + global_total: global_stats.total_ips, + queue_remaining: global_stats.queue_remaining_ips, + block_done, + block_total, + elapsed: started_at.elapsed(), + }); +} + async fn enqueue_ips( redis_url: &str, queue_name: &str, @@ -496,103 +989,6 @@ async fn redis_cli(redis_url: &str, args: &[&str]) -> Result { Ok(String::from_utf8_lossy(&output.stdout).to_string()) } -async fn init_sqlite(sqlite_path: &Path) -> Result<()> { - let sql = r#" - PRAGMA journal_mode = WAL; - PRAGMA synchronous = NORMAL; - PRAGMA busy_timeout = 5000; - CREATE TABLE IF NOT EXISTS rtsp_feed_results ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - checked_at INTEGER NOT NULL, - target_ip TEXT NOT NULL, - rtsp_path TEXT NOT NULL, - rtsp_url TEXT NOT NULL, - screenshot_path TEXT, - status TEXT NOT NULL, - error TEXT - ); - "#; - sqlite_exec(sqlite_path, sql).await -} - -async fn insert_sqlite_result( - sqlite_path: &Path, - checked_at: i64, - target_ip: &str, - rtsp_path: &str, - rtsp_url: &str, - status: String, - screenshot_path: Option<&str>, - error: Option<&str>, -) -> Result<()> { - let sql = format!( - "INSERT INTO rtsp_feed_results (checked_at, target_ip, rtsp_path, rtsp_url, screenshot_path, status, error) VALUES ({checked_at}, {target_ip}, {rtsp_path}, {rtsp_url}, {screenshot_path}, {status}, {error});", - checked_at = checked_at, - target_ip = sql_quote(target_ip), - rtsp_path = sql_quote(rtsp_path), - rtsp_url = sql_quote(rtsp_url), - screenshot_path = sql_nullable_quote(screenshot_path), - status = sql_quote(&status), - error = sql_nullable_quote(error), - ); - sqlite_exec_with_retry(sqlite_path, &sql).await -} - -async fn sqlite_exec(sqlite_path: &Path, sql: &str) -> Result<()> { - let output = timeout(Duration::from_secs(20), async { - let mut command = Command::new("sqlite3"); - command - .kill_on_drop(true) - .arg("-cmd") - .arg("PRAGMA busy_timeout = 5000;") - .arg("-cmd") - .arg("PRAGMA journal_mode = WAL;") - .arg("-cmd") - .arg("PRAGMA synchronous = NORMAL;") - .arg(sqlite_path) - .arg(sql); - command.output().await - }) - .await - .context("sqlite3 timeout")??; - - if !output.status.success() { - let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); - return Err(anyhow!( - "sqlite3 failed{}", - if stderr.is_empty() { - String::new() - } else { - format!(": {stderr}") - } - )); - } - - Ok(()) -} - -async fn sqlite_exec_with_retry(sqlite_path: &Path, sql: &str) -> Result<()> { - let mut last_err: Option = None; - - for attempt in 0..5 { - match sqlite_exec(sqlite_path, sql).await { - Ok(()) => return Ok(()), - Err(err) if is_sqlite_locked(&err) && attempt < 4 => { - last_err = Some(err); - sleep(Duration::from_millis(200 * (attempt + 1) as u64)).await; - } - Err(err) => return Err(err), - } - } - - Err(last_err.unwrap_or_else(|| anyhow!("sqlite3 failed after retries"))) -} - -fn is_sqlite_locked(err: &anyhow::Error) -> bool { - let text = err.to_string().to_ascii_lowercase(); - text.contains("database is locked") || text.contains("busy") -} - fn load_ip_list(input: &str) -> Result> { let content = match fs::read_to_string(input) { Ok(content) => content, @@ -656,25 +1052,9 @@ fn build_snapshot_path(base_dir: &Path, ip: &str, rtsp_path: &str, checked_at: i base_dir.join(format!("{ip_component}_{checked_at}_{path_component}.png")) } -fn ensure_output_dirs(snapshot_dir: &Path, sqlite_path: &Path) -> Result<()> { +fn ensure_snapshot_dir(snapshot_dir: &Path) -> Result<()> { fs::create_dir_all(snapshot_dir) - .with_context(|| format!("failed to create snapshot dir {}", snapshot_dir.display()))?; - if let Some(parent) = sqlite_path.parent() { - fs::create_dir_all(parent) - .with_context(|| format!("failed to create sqlite dir {}", parent.display()))?; - } - Ok(()) -} - -fn sql_quote(value: &str) -> String { - format!("'{}'", value.replace('\'', "''")) -} - -fn sql_nullable_quote(value: Option<&str>) -> String { - match value { - Some(value) if !value.is_empty() => sql_quote(value), - _ => "NULL".to_string(), - } + .with_context(|| format!("failed to create snapshot dir {}", snapshot_dir.display())) } fn strip_quotes(value: &str) -> &str { diff --git a/src/terminal.rs b/src/terminal.rs index 45c4456..028f98f 100644 --- a/src/terminal.rs +++ b/src/terminal.rs @@ -1,5 +1,7 @@ use std::io::{self, IsTerminal, Write}; use std::sync::Mutex; +use std::time::Duration; +use terminal_size::{terminal_size, Width}; pub const BANNER: [&str; 9] = [ " ░██████ ░██████ ", @@ -25,6 +27,23 @@ pub struct TerminalUi { lock: Mutex<()>, } +#[derive(Debug, Clone)] +pub struct RedisWorkerScreenState { + pub worker_name: String, + pub queue_name: String, + pub current_ip: String, + pub current_path: String, + pub current_status: String, + pub processed: usize, + pub captured: usize, + pub global_done: usize, + pub global_total: usize, + pub queue_remaining: Option, + pub block_done: usize, + pub block_total: usize, + pub elapsed: Duration, +} + impl TerminalUi { pub fn new(enabled: bool) -> Self { Self { @@ -56,6 +75,104 @@ impl TerminalUi { let _ = stdout.flush(); } + pub fn render_redis_worker(&self, state: &RedisWorkerScreenState) { + if !self.enabled { + return; + } + + let _guard = self.lock.lock().expect("terminal mutex poisoned"); + let mut stdout = io::stdout(); + let width = terminal_width(); + let content_width = width.saturating_sub(2).max(32); + let global_percent = percent(state.global_done, state.global_total); + let block_percent = percent(state.block_done, state.block_total); + let global_bar = progress_bar(global_percent, content_width.saturating_sub(24)); + let block_bar = progress_bar(block_percent, content_width.saturating_sub(24)); + let queue_remaining = state + .queue_remaining + .map(|value| value.to_string()) + .unwrap_or_else(|| "?".to_string()); + + let _ = write!(stdout, "\x1b[2J\x1b[H\x1b[?25l"); + let _ = writeln!( + stdout, + "{}", + fit_line( + &format!( + "CamFinder Redis Worker {} pid:{}", + state.worker_name, + std::process::id() + ), + content_width + ) + ); + let _ = writeln!(stdout, "{}", "─".repeat(content_width)); + let _ = writeln!( + stdout, + "{}", + fit_line(&format!("Fila: {}", state.queue_name), content_width) + ); + let _ = writeln!( + stdout, + "{}", + fit_line( + &format!( + "Global: [{}] {:>5.1}% {}/{} IPs redis:{}", + global_bar, + global_percent, + state.global_done, + state.global_total, + queue_remaining + ), + content_width + ) + ); + let _ = writeln!( + stdout, + "{}", + fit_line( + &format!( + "Bloco : [{}] {:>5.1}% {}/{} tentativas", + block_bar, block_percent, state.block_done, state.block_total + ), + content_width + ) + ); + let _ = writeln!(stdout, "{}", "─".repeat(content_width)); + let _ = writeln!( + stdout, + "{}", + fit_line(&format!("IP atual : {}", state.current_ip), content_width) + ); + let _ = writeln!( + stdout, + "{}", + fit_line(&format!("Path : {}", state.current_path), content_width) + ); + let _ = writeln!( + stdout, + "{}", + fit_line( + &format!("Status : {}", state.current_status), + content_width + ) + ); + let _ = writeln!( + stdout, + "{}", + fit_line( + &format!( + "Tentativas: {} Capturas: {} Tempo: {}", + state.processed, + state.captured, + format_duration(state.elapsed) + ), + content_width + ) + ); + let _ = stdout.flush(); + } + pub fn render(&self, state: &ScreenState) { if !self.enabled { return; @@ -97,3 +214,51 @@ impl TerminalUi { pub fn stdout_is_terminal() -> bool { io::stdout().is_terminal() } + +fn terminal_width() -> usize { + terminal_size() + .map(|(Width(width), _)| usize::from(width)) + .unwrap_or(80) +} + +fn percent(done: usize, total: usize) -> f64 { + if total == 0 { + 0.0 + } else { + ((done as f64 / total as f64) * 100.0).clamp(0.0, 100.0) + } +} + +fn progress_bar(percent: f64, width: usize) -> String { + let width = width.clamp(8, 80); + let filled = ((percent / 100.0) * width as f64).round() as usize; + let empty = width.saturating_sub(filled); + format!("{}{}", "█".repeat(filled), "░".repeat(empty)) +} + +fn fit_line(value: &str, width: usize) -> String { + let chars = value.chars().collect::>(); + if chars.len() <= width { + return value.to_string(); + } + if width <= 1 { + return "…".to_string(); + } + chars + .into_iter() + .take(width - 1) + .chain(std::iter::once('…')) + .collect() +} + +fn format_duration(duration: Duration) -> String { + let total = duration.as_secs(); + let hours = total / 3600; + let minutes = (total % 3600) / 60; + let seconds = total % 60; + if hours > 0 { + format!("{hours:02}:{minutes:02}:{seconds:02}") + } else { + format!("{minutes:02}:{seconds:02}") + } +}