feat: initial tools service with document scanner, image & PDF tools
Self-hosted document scanner and media processing tools. - Rust Axum gateway + worker pool with NATS JetStream - Next.js 16 frontend with shadcn/ui - Scanner pipeline: edge detection, warp, binarization, OCR - Image tools: compress, resize, convert - PDF tools: merge, split, compress - CI/CD with Docker multi-stage build Co-Authored-By: Kilo <kilo@kilo.ai>
This commit is contained in:
asepharyana
@@ -0,0 +1,186 @@
|
||||
use async_nats::Client;
|
||||
use futures::StreamExt;
|
||||
use redis::AsyncCommands;
|
||||
use uuid::Uuid;
|
||||
|
||||
use tools_common::types::{Job, JobStatus};
|
||||
|
||||
use crate::config::WorkerConfig;
|
||||
|
||||
/// NATS consumer setup and management.
|
||||
pub struct JobConsumer;
|
||||
|
||||
impl JobConsumer {
|
||||
/// Connect to NATS.
|
||||
pub async fn connect(url: &str) -> Result<Client, Box<dyn std::error::Error + Send + Sync>> {
|
||||
Ok(async_nats::connect(url).await?)
|
||||
}
|
||||
|
||||
/// Connect to Redis.
|
||||
pub async fn connect_redis(
|
||||
url: &str,
|
||||
) -> Result<redis::Client, Box<dyn std::error::Error + Send + Sync>> {
|
||||
Ok(redis::Client::open(url)?)
|
||||
}
|
||||
|
||||
/// Start consuming job messages from NATS for all tool groups.
|
||||
pub async fn start(
|
||||
nats: &Client,
|
||||
redis: &redis::Client,
|
||||
config: &WorkerConfig,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
// Subscribe to scan jobs
|
||||
let scan_sub = nats
|
||||
.queue_subscribe("tools.scan.jobs.>", "scan-workers".to_string())
|
||||
.await?;
|
||||
tracing::info!("Subscribed to tools.scan.jobs.>");
|
||||
|
||||
// Subscribe to image jobs
|
||||
let image_sub = nats
|
||||
.queue_subscribe("tools.image.jobs.>", "image-workers".to_string())
|
||||
.await?;
|
||||
tracing::info!("Subscribed to tools.image.jobs.>");
|
||||
|
||||
// Subscribe to pdf jobs
|
||||
let pdf_sub = nats
|
||||
.queue_subscribe("tools.pdf.jobs.>", "pdf-workers".to_string())
|
||||
.await?;
|
||||
tracing::info!("Subscribed to tools.pdf.jobs.>");
|
||||
|
||||
// Subscribe to cleanup scheduler
|
||||
let cleanup_sub = nats
|
||||
.subscribe("tools.scheduler.cleanup".to_string())
|
||||
.await?;
|
||||
tracing::info!("Subscribed to tools.scheduler.cleanup");
|
||||
|
||||
let redis_clone = redis.clone();
|
||||
let config_clone = config.clone();
|
||||
|
||||
// Process messages concurrently
|
||||
tokio::select! {
|
||||
_ = Self::process_subscription(scan_sub, redis.clone(), config.clone()) => {},
|
||||
_ = Self::process_subscription(image_sub, redis.clone(), config.clone()) => {},
|
||||
_ = Self::process_subscription(pdf_sub, redis.clone(), config.clone()) => {},
|
||||
_ = Self::process_cleanup(cleanup_sub, config_clone) => {},
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Process messages from a NATS subscription.
|
||||
async fn process_subscription(
|
||||
mut sub: async_nats::Subscriber,
|
||||
redis: redis::Client,
|
||||
config: WorkerConfig,
|
||||
) {
|
||||
while let Some(msg) = sub.next().await {
|
||||
if let Ok(job) = serde_json::from_slice::<Job>(&msg.payload) {
|
||||
let redis = redis.clone();
|
||||
let config = config.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let tool = job.tool.clone();
|
||||
tracing::info!(
|
||||
job_id = %job.id,
|
||||
tool = %tool.as_str(),
|
||||
"Received job"
|
||||
);
|
||||
|
||||
match Self::dispatch_job(tool, job, &redis, &config).await {
|
||||
Ok(()) => tracing::info!("Job completed successfully"),
|
||||
Err(e) => tracing::error!("Job failed: {}", e),
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Process cleanup scheduler messages.
|
||||
async fn process_cleanup(mut sub: async_nats::Subscriber, config: WorkerConfig) {
|
||||
while let Some(msg) = sub.next().await {
|
||||
tracing::info!("Running cleanup cycle");
|
||||
let redis_url = config.redis_url.clone();
|
||||
match redis::Client::open(redis_url.as_str()) {
|
||||
Ok(client) => {
|
||||
match crate::scheduler::cleanup::CleanupScheduler::run(
|
||||
&config.storage_path,
|
||||
&client,
|
||||
config.job_ttl_seconds,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(result) => {
|
||||
tracing::info!(
|
||||
"Cleanup: {} files deleted, {} bytes freed",
|
||||
result.files_deleted,
|
||||
result.bytes_freed
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("Cleanup failed: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("Failed to create Redis client for cleanup: {}", e);
|
||||
}
|
||||
}
|
||||
// Consume the message (no ack for core NATS)
|
||||
let _ = msg;
|
||||
}
|
||||
}
|
||||
|
||||
/// Dispatch a job to the appropriate handler based on tool type.
|
||||
async fn dispatch_job(
|
||||
tool: tools_common::types::Tool,
|
||||
job: Job,
|
||||
redis: &redis::Client,
|
||||
config: &WorkerConfig,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
match tool {
|
||||
tools_common::types::Tool::Scan => {
|
||||
crate::scanner::process_job(job, redis, config).await
|
||||
}
|
||||
tools_common::types::Tool::ImageCompress
|
||||
| tools_common::types::Tool::ImageResize
|
||||
| tools_common::types::Tool::ImageConvert
|
||||
| tools_common::types::Tool::RemoveBg => {
|
||||
crate::image::process_job(job, redis, config).await
|
||||
}
|
||||
tools_common::types::Tool::PdfMerge
|
||||
| tools_common::types::Tool::PdfSplit
|
||||
| tools_common::types::Tool::ImagesToPdf
|
||||
| tools_common::types::Tool::PdfCompress
|
||||
| tools_common::types::Tool::PdfToImages => {
|
||||
crate::pdf::process_job(job, redis, config).await
|
||||
}
|
||||
_ => {
|
||||
tracing::warn!(tool = %tool.as_str(), "Tool handler not yet implemented");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Update job result in Redis after processing.
|
||||
pub async fn update_job_result(
|
||||
conn: &mut impl AsyncCommands,
|
||||
job_id: Uuid,
|
||||
result_path: &str,
|
||||
ttl_seconds: u64,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
let key = format!("job:{}", job_id);
|
||||
let json: String = conn
|
||||
.get(&key)
|
||||
.await
|
||||
.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { Box::new(e) })?;
|
||||
let mut job: Job = serde_json::from_str(&json)?;
|
||||
job.status = JobStatus::Completed;
|
||||
job.result_path = Some(result_path.to_string());
|
||||
let updated = serde_json::to_string(&job)?;
|
||||
let _: () = conn
|
||||
.set_ex(key, updated, ttl_seconds)
|
||||
.await
|
||||
.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { Box::new(e) })?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod consumer;
|
||||
pub mod progress;
|
||||
@@ -0,0 +1,82 @@
|
||||
use redis::AsyncCommands;
|
||||
use uuid::Uuid;
|
||||
|
||||
use tools_common::types::{JobStatus, Tool};
|
||||
|
||||
/// Reports progress from worker to NATS and Redis.
|
||||
pub struct ProgressReporter {
|
||||
redis: redis::Client,
|
||||
nats: async_nats::Client,
|
||||
job_id: Uuid,
|
||||
tool: Tool,
|
||||
}
|
||||
|
||||
impl ProgressReporter {
|
||||
pub fn new(redis: redis::Client, nats: async_nats::Client, job_id: Uuid, tool: Tool) -> Self {
|
||||
Self {
|
||||
redis,
|
||||
nats,
|
||||
job_id,
|
||||
tool,
|
||||
}
|
||||
}
|
||||
|
||||
/// Report progress: updates Redis and publishes to NATS.
|
||||
pub async fn report(
|
||||
&self,
|
||||
status: JobStatus,
|
||||
stage: &str,
|
||||
progress: u8,
|
||||
message: &str,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
// Update Redis
|
||||
if let Ok(mut conn) = self.redis.get_multiplexed_async_connection().await {
|
||||
let key = format!("job:{}", self.job_id);
|
||||
if let Ok(json) = conn.get::<_, String>(&key).await {
|
||||
if let Ok(mut job) = serde_json::from_str::<tools_common::types::Job>(&json) {
|
||||
job.status = status.clone();
|
||||
let updated = serde_json::to_string(&job).unwrap_or(json);
|
||||
let _: Result<(), _> = conn.set_ex(key, updated, job.ttl_seconds).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Publish to NATS
|
||||
let progress_msg = tools_common::types::JobProgress {
|
||||
job_id: self.job_id,
|
||||
status,
|
||||
stage: stage.to_string(),
|
||||
progress,
|
||||
message: message.to_string(),
|
||||
};
|
||||
|
||||
let subject = format!("tools.{}.progress.{}", self.tool.subject_prefix(), self.job_id);
|
||||
if let Ok(payload) = serde_json::to_vec(&progress_msg) {
|
||||
let _ = self.nats.publish(subject, payload.into()).await;
|
||||
}
|
||||
|
||||
tracing::debug!(
|
||||
job_id = %self.job_id,
|
||||
stage = %stage,
|
||||
progress = %progress,
|
||||
"Progress update"
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn job_id(&self) -> Uuid {
|
||||
self.job_id
|
||||
}
|
||||
}
|
||||
|
||||
impl Clone for ProgressReporter {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
redis: self.redis.clone(),
|
||||
nats: self.nats.clone(),
|
||||
job_id: self.job_id,
|
||||
tool: self.tool.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user