feat: Add database migration service with Alembic for PostgreSQL, integrating job management and counting functionalities. Update Docker Compose to include db-migrations service, enhancing the microservices architecture. Revise documentation for clarity on database schema and migration processes.
This commit is contained in:
erik
@@ -1,51 +1,59 @@
|
||||
# syntax=docker/dockerfile:1
|
||||
|
||||
FROM golang:1.23-bookworm AS builder
|
||||
# ── Etapa 1: Descargar ONNX Runtime ───────────────────────────────────────────
|
||||
FROM debian:bookworm-slim AS onnx-download
|
||||
|
||||
ARG ONNX_VERSION=1.22.0
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends wget ca-certificates \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
RUN wget -q https://github.com/microsoft/onnxruntime/releases/download/v${ONNX_VERSION}/onnxruntime-linux-x64-${ONNX_VERSION}.tgz \
|
||||
&& tar -xzf onnxruntime-linux-x64-${ONNX_VERSION}.tgz \
|
||||
&& mkdir -p /onnx/lib \
|
||||
&& cp onnxruntime-linux-x64-${ONNX_VERSION}/lib/libonnxruntime.so* /onnx/lib/ \
|
||||
&& rm -rf onnxruntime-linux-x64-${ONNX_VERSION}*
|
||||
|
||||
# ── Etapa 2: Build ─────────────────────────────────────────────────────────────
|
||||
FROM golang:1.25-bookworm AS builder
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
gcc g++ pkg-config \
|
||||
libtesseract-dev libleptonica-dev \
|
||||
tesseract-ocr tesseract-ocr-spa \
|
||||
wget ca-certificates \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# ONNX Runtime
|
||||
ARG ONNX_VERSION=1.22.0
|
||||
RUN wget -q https://github.com/microsoft/onnxruntime/releases/download/v${ONNX_VERSION}/onnxruntime-linux-x64-${ONNX_VERSION}.tgz \
|
||||
&& tar -xzf onnxruntime-linux-x64-${ONNX_VERSION}.tgz \
|
||||
&& cp onnxruntime-linux-x64-${ONNX_VERSION}/lib/libonnxruntime.so* /usr/local/lib/ \
|
||||
&& ldconfig
|
||||
COPY --from=onnx-download /onnx/lib/ /usr/local/lib/
|
||||
RUN ldconfig
|
||||
|
||||
ENV CGO_ENABLED=1
|
||||
ENV ONNXRUNTIME_SHARED_LIBRARY=/usr/local/lib/libonnxruntime.so
|
||||
# Limitar paralelismo CGO para evitar OOM en Docker Desktop
|
||||
ENV GOMAXPROCS=2
|
||||
|
||||
WORKDIR /src
|
||||
COPY go.mod go.sum ./
|
||||
RUN go mod download
|
||||
|
||||
COPY . .
|
||||
|
||||
RUN go build -o /out/api ./cmd/api && go build -o /out/worker ./cmd/worker
|
||||
# -p 2 limita los paquetes compilados en paralelo (menor uso de RAM)
|
||||
RUN go build -p 2 -o /out/api ./cmd/api
|
||||
|
||||
# ── Etapa 3: Runtime ───────────────────────────────────────────────────────────
|
||||
FROM debian:bookworm-slim AS runtime
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
ca-certificates curl wget \
|
||||
ca-certificates curl \
|
||||
tesseract-ocr tesseract-ocr-spa \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
ARG ONNX_VERSION=1.22.0
|
||||
RUN wget -q https://github.com/microsoft/onnxruntime/releases/download/v${ONNX_VERSION}/onnxruntime-linux-x64-${ONNX_VERSION}.tgz \
|
||||
&& tar -xzf onnxruntime-linux-x64-${ONNX_VERSION}.tgz \
|
||||
&& cp onnxruntime-linux-x64-${ONNX_VERSION}/lib/libonnxruntime.so* /usr/local/lib/ \
|
||||
&& rm -rf onnxruntime-linux-x64-${ONNX_VERSION}* \
|
||||
&& ldconfig
|
||||
COPY --from=onnx-download /onnx/lib/ /usr/local/lib/
|
||||
RUN ldconfig
|
||||
|
||||
ENV ONNXRUNTIME_SHARED_LIBRARY=/usr/local/lib/libonnxruntime.so
|
||||
|
||||
WORKDIR /app
|
||||
COPY --from=builder /out/api /app/api
|
||||
COPY --from=builder /out/worker /app/worker
|
||||
COPY api/openapi.yaml /app/openapi.yaml
|
||||
COPY --from=builder /out/api /app/api
|
||||
COPY api/openapi.yaml /app/openapi.yaml
|
||||
|
||||
ENV OPENAPI_PATH=/app/openapi.yaml
|
||||
ENV MODEL_PATH=/app/models/model.onnx
|
||||
|
||||
@@ -1,14 +0,0 @@
|
||||
FROM golang:1.23-bookworm AS builder
|
||||
|
||||
WORKDIR /src
|
||||
COPY go.mod go.sum ./
|
||||
RUN go mod download
|
||||
COPY . .
|
||||
RUN go build -o /out/autoscaler ./cmd/autoscaler
|
||||
|
||||
FROM docker:27-cli
|
||||
|
||||
COPY --from=builder /out/autoscaler /usr/local/bin/autoscaler
|
||||
|
||||
WORKDIR /project
|
||||
ENTRYPOINT ["/usr/local/bin/autoscaler"]
|
||||
@@ -1,339 +1,175 @@
|
||||
# Services Count — API de Facturas Médicas
|
||||
# count-api
|
||||
|
||||
|
||||
Microservicio Go que cuenta servicios facturados en facturas médicas PDF.
|
||||
Aplica OCR (Tesseract) + clustering espacial (DBSCAN) + clasificación ML (Random Forest vía ONNX Runtime).
|
||||
|
||||
Microservicio de inferencia en Go que cuenta servicios facturados en facturas médicas PDF. El entrenamiento del modelo vive en el repositorio hermano [services-count-facturas-medicas-ml](https://github.com/builker-col/services-count-facturas-medicas-ml); este repositorio despliega la API HTTP y los workers de procesamiento.
|
||||
Es un servicio **sin estado**: recibe PDFs por HTTP, procesa en memoria y devuelve el resultado. No persiste datos ni depende de bases de datos, colas ni almacenamiento de objetos.
|
||||
|
||||
## Características
|
||||
La persistencia de resultados la maneja el servicio `extraction` (Python) a través de los endpoints `/count-jobs`.
|
||||
|
||||
- Conteo **síncrono** (1–50 PDFs por petición) y **asíncrono** (hasta 50 000 PDFs por job).
|
||||
- Pipeline de inferencia: OCR → clustering DBSCAN → Random Forest vía ONNX.
|
||||
- Escalado horizontal de workers con Docker Compose.
|
||||
- Especificación OpenAPI y descriptor de herramientas para agentes AI.
|
||||
---
|
||||
|
||||
## Arquitectura
|
||||
|
||||
El sistema se compone de una API HTTP, workers consumidores y tres servicios de infraestructura (Redis, NATS JetStream y MinIO), definidos en [`docker-compose.yml`](docker-compose.yml).
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
Client[Cliente_o_Agente]
|
||||
API[API_Go_8080]
|
||||
Worker[Worker_Go_xN]
|
||||
Redis[(Redis)]
|
||||
NATS[NATS_JetStream]
|
||||
MinIO[(MinIO)]
|
||||
Pipeline[OCR_Tesseract_ONNX]
|
||||
|
||||
Client -->|sync_POST_v1_count| API
|
||||
Client -->|async_POST_v1_jobs| API
|
||||
API --> Pipeline
|
||||
API --> Redis
|
||||
API --> MinIO
|
||||
API -->|publish| NATS
|
||||
NATS -->|consume| Worker
|
||||
Worker --> MinIO
|
||||
Worker --> Pipeline
|
||||
Worker --> Redis
|
||||
```
|
||||
|
||||
### Pipeline de inferencia
|
||||
|
||||
Por cada PDF se procesa la primera página siguiendo este flujo (`internal/pipeline`):
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
PDF[PDF_primera_pagina] --> Render[Render_MuPDF]
|
||||
Render --> OCR[Tesseract_spa]
|
||||
OCR --> Cluster[DBSCAN_filas]
|
||||
Cluster --> Features[8_features]
|
||||
Features --> Infer[ONNX_RandomForest]
|
||||
Infer --> Result[service_count_y_services]
|
||||
```
|
||||
|
||||
1. **Render** — La primera página se rasteriza con MuPDF (go-fitz) a 300 DPI.
|
||||
2. **OCR** — Tesseract extrae cajas de texto en español.
|
||||
3. **Clustering** — DBSCAN agrupa palabras en filas candidatas.
|
||||
4. **Features** — Se calculan 8 atributos por fila (códigos, moneda, posición, etc.).
|
||||
5. **Inferencia** — Un Random Forest exportado a ONNX clasifica cada fila como servicio o no.
|
||||
|
||||
## Componentes
|
||||
|
||||
| Componente | Binario / servicio | Paquete | Responsabilidad |
|
||||
|------------|-------------------|---------|-----------------|
|
||||
| API HTTP | `cmd/api` | `internal/api` | Endpoints REST, gestión de jobs, health probes |
|
||||
| Worker | `cmd/worker` | `internal/worker` | Consume NATS, descarga PDFs de MinIO, persiste resultados |
|
||||
| Pipeline | — | `internal/pipeline` | Orquestación OCR → clustering → features → inferencia |
|
||||
| Modelo ONNX | — | `internal/model` | Carga `model.onnx` e `inference.json` |
|
||||
| OCR | — | `internal/ocr` | Render PDF (go-fitz) + Tesseract (gosseract) |
|
||||
| Estado de jobs | Redis 7 | `internal/store` | Metadatos, resultados y progreso (TTL 72 h) |
|
||||
| Cola | NATS JetStream | `internal/queue` | Stream `PDF_JOBS`, subject `pdf.process.*` |
|
||||
| Object storage | MinIO | `internal/storage` | Almacenamiento de PDFs en jobs asíncronos |
|
||||
|
||||
### Estructura del repositorio
|
||||
|
||||
```
|
||||
├── cmd/
|
||||
│ ├── api/ # Punto de entrada HTTP
|
||||
│ └── worker/ # Punto de entrada worker
|
||||
├── internal/
|
||||
│ ├── api/ # Servidor HTTP (chi)
|
||||
│ ├── config/ # Variables de entorno
|
||||
│ ├── domain/ # Tipos de dominio
|
||||
│ ├── model/ # Inferencia ONNX
|
||||
│ ├── ocr/ # Extracción OCR
|
||||
│ ├── pipeline/ # Procesamiento por PDF
|
||||
│ ├── queue/ # NATS JetStream
|
||||
│ ├── storage/ # Cliente MinIO
|
||||
│ ├── store/ # Estado en Redis
|
||||
│ └── worker/ # Runner del worker
|
||||
├── api/openapi.yaml # Especificación OpenAPI 3.1
|
||||
├── mcp/tool.json # Descriptor de herramientas para agentes
|
||||
├── web/ # Cliente demo (embebido en la API en /demo/)
|
||||
├── scripts/ # Export ONNX, paridad con Python
|
||||
└── models/latest/ # model.onnx + inference.json (montado en Docker)
|
||||
```
|
||||
|
||||
## Modos de operación
|
||||
|
||||
| Modo | Endpoint principal | Límite | Cuándo usarlo |
|
||||
|------|-------------------|--------|---------------|
|
||||
| Síncrono | `POST /v1/count` | 50 PDFs | Pruebas, integraciones en tiempo real, lotes pequeños |
|
||||
| Asíncrono | `POST /v1/jobs` + poll | 50 000 PDFs/job | MinIO + NATS + workers (path producción) |
|
||||
|
||||
**Flujo asíncrono (producción):**
|
||||
|
||||
1. Crear job (`POST /v1/jobs`) — shell vacío, prefijo MinIO o lista de claves.
|
||||
2. Subir PDFs por chunks (`POST /v1/jobs/{jobID}/upload`, máx. 500 por chunk) → **MinIO Put** → **NATS Publish**.
|
||||
3. Workers consumen la cola, descargan de MinIO, OCR, guardan en Redis.
|
||||
4. Consultar progreso (`GET /v1/jobs/{jobID}?wait_ms=30000` para long-polling).
|
||||
5. Obtener resultados paginados (`GET /v1/jobs/{jobID}/results?cursor=&limit=`).
|
||||
|
||||
## API REST
|
||||
|
||||
Especificación completa: `GET /openapi.json` o [`api/openapi.yaml`](api/openapi.yaml).
|
||||
## Endpoints
|
||||
|
||||
| Método | Ruta | Descripción |
|
||||
|--------|------|-------------|
|
||||
| `GET` | `/health` | Liveness probe |
|
||||
| `GET` | `/ready` | Readiness probe |
|
||||
| `POST` | `/v1/count` | Conteo síncrono (multipart `files`) |
|
||||
| `POST` | `/v1/jobs` | Crear job asíncrono |
|
||||
| `POST` | `/v1/jobs/{jobID}/upload` | Subir chunk de PDFs |
|
||||
| `GET` | `/v1/jobs/{jobID}` | Estado y progreso del job |
|
||||
| `GET` | `/v1/jobs/{jobID}/results` | Resultados paginados |
|
||||
| `DELETE` | `/v1/jobs/{jobID}` | Cancelar job |
|
||||
| `GET` | `/health` | Liveness — `{"status":"ok"}` |
|
||||
| `GET` | `/ready` | Readiness — `{"status":"ready"}` |
|
||||
| `GET` | `/openapi.json` | Especificación OpenAPI |
|
||||
| `POST` | `/v1/count` | Conteo síncrono de PDFs |
|
||||
|
||||
**Parámetros relevantes:**
|
||||
### POST /v1/count
|
||||
|
||||
- `include_rejected` (query) — Si es `true`, incluye todas las filas OCR con el campo `is_service`.
|
||||
- Jobs asíncronos devuelven `202 Accepted`; la cola rechaza con `429` si supera `QUEUE_MAX_PENDING`.
|
||||
**Request**
|
||||
|
||||
**Respuesta por PDF:**
|
||||
```
|
||||
Content-Type: multipart/form-data
|
||||
Query: include_rejected=true|false (default: false)
|
||||
Field: files — uno o más archivos PDF (máx. SYNC_MAX_FILES)
|
||||
```
|
||||
|
||||
**Response 200**
|
||||
|
||||
```json
|
||||
{
|
||||
"filename": "factura.pdf",
|
||||
"service_count": 2,
|
||||
"status": "ok",
|
||||
"services": [
|
||||
{ "text": "890201 Consulta ...", "probability": 0.94 }
|
||||
]
|
||||
"results": [
|
||||
{
|
||||
"filename": "factura_001.pdf",
|
||||
"service_count": 3,
|
||||
"status": "ok",
|
||||
"services": [
|
||||
{ "text": "CONSULTA MEDICA", "probability": 0.91 },
|
||||
{ "text": "LABORATORIO HEMOGRAMA", "probability": 0.87 },
|
||||
{ "text": "RADIOGRAFIA TORAX", "probability": 0.83 }
|
||||
]
|
||||
}
|
||||
],
|
||||
"processed": 1,
|
||||
"failed": 0
|
||||
}
|
||||
```
|
||||
|
||||
## Stack tecnológico
|
||||
**Response 4xx/5xx**
|
||||
|
||||
| Área | Tecnología |
|
||||
|------|------------|
|
||||
| Lenguaje | Go 1.23 |
|
||||
| HTTP router | [go-chi/chi](https://github.com/go-chi/chi) v5 |
|
||||
| OCR | go-fitz (MuPDF) + gosseract (Tesseract) |
|
||||
| ML | ONNX Runtime + Random Forest exportado |
|
||||
| Cola | NATS JetStream |
|
||||
| Estado | Redis 7 |
|
||||
| Storage | MinIO |
|
||||
| Contenedores | Docker multi-stage ([`Dockerfile`](Dockerfile)) |
|
||||
|
||||
## Inicio rápido
|
||||
|
||||
### 1. Exportar el modelo ONNX
|
||||
|
||||
Desde el repositorio de entrenamiento, generar los artefactos de inferencia:
|
||||
|
||||
```bash
|
||||
pip install -r requirements-export.txt
|
||||
python scripts/export_onnx.py \
|
||||
--joblib ../services-count/models/latest/model.joblib \
|
||||
--metadata ../services-count/models/v1.0.0/model_metadata.json
|
||||
```json
|
||||
{ "error": "descripción del error" }
|
||||
```
|
||||
|
||||
Esto produce `models/latest/model.onnx` e `inference.json`.
|
||||
---
|
||||
|
||||
### 2. Levantar con Docker Compose
|
||||
## Pipeline de procesamiento
|
||||
|
||||
```bash
|
||||
docker compose up --build
|
||||
```
|
||||
PDF bytes (memoria)
|
||||
│
|
||||
▼ gen2brain/go-fitz (MuPDF)
|
||||
Render de la primera página a imagen (OCR_DPI, default 300)
|
||||
│
|
||||
▼ otiai10/gosseract (Tesseract)
|
||||
OCR palabra a palabra con coordenadas (X, Y, W, H)
|
||||
│
|
||||
▼ DBSCAN 1D sobre coordenada Y
|
||||
Agrupación de palabras en filas candidatas
|
||||
│
|
||||
▼ 8 features por fila
|
||||
first_token_is_code · currency_count · integer_count · x_span_norm
|
||||
y_norm · token_count · special_char_ratio · digit_ratio
|
||||
│
|
||||
▼ ONNX Runtime (Random Forest)
|
||||
P(servicio) por fila → filtrar por threshold (inference.json)
|
||||
│ Si el modelo no está disponible → heurística: todas las filas = servicio
|
||||
▼
|
||||
PDFResult { service_count, services[], status }
|
||||
```
|
||||
|
||||
| Servicio | Puerto / acceso |
|
||||
|----------|-----------------|
|
||||
| API | `http://localhost:8080` |
|
||||
| MinIO Console | `http://localhost:9002` |
|
||||
| Workers | 1 réplica inicial; autoscaler sube hasta 8 según cola |
|
||||
| Autoscaler | Escala `worker` vía Docker socket según `PendingCount` en Redis |
|
||||
---
|
||||
|
||||
### 3. Probar
|
||||
## Modelo ONNX
|
||||
|
||||
```bash
|
||||
# Síncrono
|
||||
curl -F "files=@factura.pdf" http://localhost:8080/v1/count
|
||||
El modelo es opcional. Si no está disponible, el servicio arranca con una advertencia y usa heurística de conteo.
|
||||
|
||||
# Asíncrono (PDFs ya en MinIO bajo un prefijo)
|
||||
curl -X POST http://localhost:8080/v1/jobs \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"prefix":"jobs/batch-1/"}'
|
||||
Los archivos deben montarse en `/app/models/`:
|
||||
|
||||
```
|
||||
/app/models/model.onnx — modelo Random Forest exportado con skl2onnx
|
||||
/app/models/inference.json — threshold y feature_columns
|
||||
```
|
||||
|
||||
## Cliente web de ejemplo
|
||||
|
||||
Con la API en ejecución, abre **[http://localhost:8080/demo/](http://localhost:8080/demo/)** para probar la integración desde el navegador. Los archivos estáticos se sirven desde la propia API (mismo origen, sin CORS).
|
||||
|
||||
| Pestaña | Demuestra |
|
||||
|---------|-----------|
|
||||
| **Síncrono** | `POST /v1/count` con subida multipart de hasta 50 PDFs |
|
||||
| **Asíncrono** | `POST /v1/jobs` → upload → polling de estado → resultados paginados |
|
||||
|
||||
El modo asíncrono requiere que workers, Redis, NATS y MinIO estén levantados (igual que con `docker compose up`). Código fuente en [`web/`](web/).
|
||||
|
||||
## Desarrollo local
|
||||
|
||||
Requisitos adicionales a Go 1.23+:
|
||||
|
||||
- Tesseract con datos de idioma `spa`
|
||||
- ONNX Runtime (la variable `ONNXRUNTIME_SHARED_LIBRARY` apunta a la librería compartida)
|
||||
- Redis, NATS y MinIO en ejecución
|
||||
Para generar el modelo desde un `model.joblib` existente:
|
||||
|
||||
```bash
|
||||
go build -o api ./cmd/api
|
||||
go build -o worker ./cmd/worker
|
||||
cd scripts
|
||||
pip install -r ../requirements-export.txt
|
||||
python export_onnx.py --joblib /ruta/model.joblib --out ../models/latest
|
||||
```
|
||||
|
||||
Configurar las variables de entorno para apuntar a los servicios locales (ver tabla siguiente).
|
||||
---
|
||||
|
||||
## Configuración
|
||||
|
||||
Variables definidas en [`internal/config/config.go`](internal/config/config.go):
|
||||
## Variables de entorno
|
||||
|
||||
| Variable | Default | Descripción |
|
||||
|----------|---------|-------------|
|
||||
| `PORT` | `8080` | Puerto de la API |
|
||||
| `REDIS_URL` | `redis://redis:6379/0` | Conexión Redis |
|
||||
| `NATS_URL` | `nats://nats:4222` | Conexión NATS |
|
||||
| `NATS_STREAM` | `PDF_JOBS` | Nombre del stream JetStream |
|
||||
| `NATS_SUBJECT_PREFIX` | `pdf.process` | Prefijo de subjects |
|
||||
| `MINIO_ENDPOINT` | `minio:9000` | Endpoint MinIO |
|
||||
| `MINIO_ACCESS_KEY` | `minioadmin` | Access key |
|
||||
| `MINIO_SECRET_KEY` | `minioadmin` | Secret key |
|
||||
| `MINIO_BUCKET` | `pdf-count` | Bucket de PDFs |
|
||||
| `MINIO_USE_SSL` | `false` | Usar SSL con MinIO |
|
||||
| `PORT` | `8080` | Puerto HTTP |
|
||||
| `MODEL_PATH` | `/app/models/model.onnx` | Ruta al modelo ONNX |
|
||||
| `INFERENCE_CONFIG` | `/app/models/inference.json` | Metadata de inferencia |
|
||||
| `OCR_DPI` | `300` | DPI de renderizado PDF (150 reduce latencia ~2–4×) |
|
||||
| `INFERENCE_CONFIG` | `/app/models/inference.json` | Threshold y metadatos del modelo |
|
||||
| `OPENAPI_PATH` | `api/openapi.yaml` | Ruta a la especificación OpenAPI |
|
||||
| `OCR_DPI` | `300` | DPI de renderizado PDF para OCR |
|
||||
| `OCR_LANG` | `spa` | Idioma Tesseract |
|
||||
| `DBSCAN_EPS` | `10` | Parámetro epsilon de DBSCAN |
|
||||
| `DBSCAN_MIN_SAMPLES` | `1` | Mínimo de muestras DBSCAN |
|
||||
| `SYNC_MAX_FILES` | `50` | Máximo de PDFs en modo síncrono |
|
||||
| `JOB_MAX_FILES` | `50000` | Máximo de PDFs por job |
|
||||
| `QUEUE_MAX_PENDING` | `100000` | Límite de mensajes pendientes en cola |
|
||||
| `WORKER_CONCURRENCY` | `4` | Goroutines paralelas por worker |
|
||||
| `MAX_PDF_SIZE_MB` | `20` | Tamaño máximo por PDF |
|
||||
| `PDF_PROCESS_TIMEOUT_SEC` | `120` | Timeout de procesamiento por PDF |
|
||||
| `JOB_RESULT_TTL_HOURS` | `72` | TTL de resultados en Redis |
|
||||
| `OCR_POOL_SIZE` | `4` (API) / `WORKER_CONCURRENCY` (worker) | Clientes Tesseract reutilizables en pool |
|
||||
| `SYNC_PROCESS_CONCURRENCY` | `4` (o `OCR_POOL_SIZE`) | PDFs procesados en paralelo en `/v1/count` |
|
||||
| `JOB_INLINE_PROCESSING` | `false` | `true` procesa uploads en la API sin MinIO (solo dev local) |
|
||||
| `LOG_PIPELINE_TIMING` | `false` | Logs de duración por etapa del pipeline |
|
||||
| `NATS_FETCH_MAX_WAIT_MS` | `300` | Tiempo máximo de espera del worker al leer NATS |
|
||||
| `OCR_POOL_SIZE` | `0` (→ 4) | Tamaño del pool de clientes Tesseract |
|
||||
| `DBSCAN_EPS` | `10` | Épsilon de clustering DBSCAN (píxeles en Y) |
|
||||
| `DBSCAN_MIN_SAMPLES` | `1` | Mínimo de puntos por cluster DBSCAN |
|
||||
| `SYNC_MAX_FILES` | `50` | Máximo de PDFs por request a `/v1/count` |
|
||||
| `SYNC_PROCESS_CONCURRENCY` | `0` (→ OCR_POOL_SIZE) | Goroutines paralelas por request |
|
||||
| `MAX_PDF_SIZE_MB` | `20` | Tamaño máximo permitido por PDF |
|
||||
| `PDF_PROCESS_TIMEOUT_SEC` | `120` | Timeout por PDF |
|
||||
| `LOG_PIPELINE_TIMING` | `false` | Loguear tiempos de cada etapa |
|
||||
|
||||
## Rendimiento
|
||||
---
|
||||
|
||||
El cuello de botella de `/v1/count` es el **OCR con Tesseract**, no la inferencia ONNX. Los clientes Tesseract se crean una vez al arranque y se reutilizan desde un pool (`internal/ocr/pool.go`).
|
||||
## Dependencias en producción (Docker)
|
||||
|
||||
| Etapa | Tiempo típico (1 PDF) |
|
||||
|-------|------------------------|
|
||||
| Render PDF (MuPDF) | ~100–300 ms |
|
||||
| OCR Tesseract | ~3–8 s |
|
||||
| Clustering + ONNX | ~10–50 ms |
|
||||
- `libonnxruntime.so` — ONNX Runtime (copiado en build multi-etapa)
|
||||
- `tesseract-ocr tesseract-ocr-spa` — OCR en español
|
||||
- `libmupdf` — via go-fitz, renderizado de PDFs
|
||||
|
||||
**Latencia esperada (síncrono):** ~3–10 s por PDF. Varios PDFs se procesan en paralelo (`SYNC_PROCESS_CONCURRENCY`, default 4). Ejemplo: 20 PDFs ≈ `ceil(20/4) × 2.85s ≈ 15s`.
|
||||
Sin Redis, MinIO, NATS ni bases de datos.
|
||||
|
||||
**Latencia esperada (asíncrono por upload):** path producción MinIO → NATS → worker. Con `WORKER_CONCURRENCY=4` y autoscaler activo, 16 PDFs ≈ `ceil(16/4) × 2.85s` más overhead Put/Get (~15–25 s). 100 PDFs escala hasta 8 workers (32 slots OCR).
|
||||
---
|
||||
|
||||
**Asíncrono por prefix MinIO** (`POST /v1/jobs {"prefix":"..."}`): mismos workers; los PDFs ya están en el bucket.
|
||||
|
||||
**Ajustes para reducir latencia:**
|
||||
|
||||
- `OCR_DPI=150` — reduce tiempo de render y OCR (~2–4× más rápido) con posible ligera pérdida de precisión. Validar con `scripts/compare_parity.py` antes de usar en producción.
|
||||
- `OCR_POOL_SIZE=4` y `WORKER_CONCURRENCY=4` — paralelismo por réplica worker.
|
||||
- Upload pipeline publica en NATS tras cada Put (no al final del lote).
|
||||
- `LOG_PIPELINE_TIMING=true` — logs por etapa: `ocr_render`, `ocr_tesseract`, `cluster`, `infer`, `total`.
|
||||
- `NATS_FETCH_MAX_WAIT_MS=300` — el worker recoge mensajes con menor espera.
|
||||
- `GET /v1/jobs/{id}?wait_ms=30000` — long-polling en lugar de poll cada 500 ms.
|
||||
|
||||
## Escalado
|
||||
|
||||
### Autoscaler (docker compose)
|
||||
|
||||
El servicio `autoscaler` lee `PendingCount` en Redis cada 5 s y ajusta réplicas `worker`:
|
||||
## Estructura del código
|
||||
|
||||
```
|
||||
desired = clamp(ceil(pending / WORKER_SLOTS), MIN_WORKERS, MAX_WORKERS)
|
||||
cmd/api/ — punto de entrada del servidor HTTP
|
||||
internal/
|
||||
api/ — router y handler countSync
|
||||
config/ — carga de variables de entorno
|
||||
domain/ — tipos PDFResult, CountResponse, ServiceRow
|
||||
model/ — sesión ONNX Runtime (Predictor)
|
||||
ocr/ — pool de clientes Tesseract + renderizado MuPDF
|
||||
pipeline/ — orquestación OCR → DBSCAN → features → ML
|
||||
processor.go — ProcessPDF por documento
|
||||
clustering.go — DBSCAN 1D en Y
|
||||
features.go — 8 features por fila
|
||||
api/openapi.yaml — especificación OpenAPI 3.1
|
||||
models/latest/ — model.onnx + inference.json (gitignored, montar en Docker)
|
||||
scripts/ — export_onnx.py para generar el modelo
|
||||
Dockerfile — build multi-etapa: ONNX download + Go builder + runtime slim
|
||||
```
|
||||
|
||||
| Variable | Default | Descripción |
|
||||
|----------|---------|-------------|
|
||||
| `MIN_WORKERS` | `1` | Réplicas mínimas |
|
||||
| `MAX_WORKERS` | `8` | Réplicas máximas |
|
||||
| `WORKER_SLOTS` | `4` | PDFs en paralelo por réplica (`WORKER_CONCURRENCY`) |
|
||||
| `SCALE_DOWN_IDLE_SEC` | `60` | Segundos sin cola antes de bajar 1 réplica |
|
||||
---
|
||||
|
||||
Requiere montar `/var/run/docker.sock` (solo entorno local). En Kubernetes usar HPA sobre lag de NATS o métricas custom.
|
||||
|
||||
### Escalado manual
|
||||
## Arrancar localmente (Docker)
|
||||
|
||||
```bash
|
||||
docker compose up --scale worker=8 -d
|
||||
# Desde la raíz del repositorio:
|
||||
docker compose up --build count-api
|
||||
```
|
||||
|
||||
Cada worker procesa hasta `WORKER_CONCURRENCY` PDFs en paralelo. NATS JetStream distribuye los mensajes entre réplicas con política work-queue.
|
||||
El servicio queda disponible en `http://localhost:8081`.
|
||||
|
||||
## Integración con agentes (MCP)
|
||||
|
||||
[`mcp/tool.json`](mcp/tool.json) es un **descriptor declarativo** de herramientas para agentes AI (Cursor, etc.). No incluye un servidor MCP embebido.
|
||||
|
||||
| Herramienta | Modo | Flujo |
|
||||
|-------------|------|-------|
|
||||
| `count_pdf_services` | Síncrono | `POST /v1/count` con multipart |
|
||||
| `count_pdf_services_batch` | Asíncrono | Crear job → poll estado → resultados paginados |
|
||||
|
||||
## Paridad con Python
|
||||
|
||||
Scripts para validar que la inferencia Go reproduce los resultados del pipeline Python original:
|
||||
Para probar el endpoint:
|
||||
|
||||
```bash
|
||||
set SERVICES_COUNT_ROOT=../services-count
|
||||
python scripts/parity_baseline.py
|
||||
python scripts/compare_parity.py
|
||||
curl -X POST http://localhost:8081/v1/count \
|
||||
-F "files=@factura.pdf" | jq .
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
| Síntoma | Causa probable | Acción |
|
||||
|---------|------------------|--------|
|
||||
| API no arranca | Modelo ONNX ausente | Ejecutar `scripts/export_onnx.py` y verificar `models/latest/` |
|
||||
| `/ready` falla | Redis o NATS no disponibles | Comprobar healthchecks en `docker compose ps` |
|
||||
| Jobs en `pending` sin avance | Workers no consumen cola | Verificar logs del worker y conectividad a NATS/MinIO |
|
||||
| Job atascado en `running` con N/N procesados | Race en metadata del job (versiones antiguas) | Reconstruir con `docker compose up --build`; `GetJob` ahora finaliza jobs completados |
|
||||
| `429` al crear job | Cola saturada | Escalar workers o esperar a que baje `QUEUE_MAX_PENDING` |
|
||||
| Resultados vacíos o erróneos | OCR o modelo desactualizado | Revisar `OCR_LANG=spa`, reexportar ONNX desde el repo ML |
|
||||
|
||||
@@ -6,15 +6,11 @@ import (
|
||||
"os"
|
||||
"os/signal"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/pys/pdf-count/internal/api"
|
||||
"github.com/pys/pdf-count/internal/config"
|
||||
"github.com/pys/pdf-count/internal/model"
|
||||
"github.com/pys/pdf-count/internal/ocr"
|
||||
"github.com/pys/pdf-count/internal/queue"
|
||||
"github.com/pys/pdf-count/internal/storage"
|
||||
"github.com/pys/pdf-count/internal/store"
|
||||
)
|
||||
|
||||
func main() {
|
||||
@@ -24,9 +20,12 @@ func main() {
|
||||
|
||||
predictor, err := model.NewPredictor(cfg.ModelPath, cfg.InferencePath)
|
||||
if err != nil {
|
||||
log.Fatalf("model: %v", err)
|
||||
log.Printf("WARN: modelo ML no disponible (%v) — usando heuristica de conteo", err)
|
||||
predictor = nil
|
||||
} else {
|
||||
defer predictor.Close()
|
||||
log.Printf("modelo ONNX cargado: %s", cfg.ModelPath)
|
||||
}
|
||||
defer predictor.Close()
|
||||
|
||||
poolSize := cfg.OCRPoolSize
|
||||
if poolSize <= 0 {
|
||||
@@ -38,37 +37,16 @@ func main() {
|
||||
}
|
||||
defer ocrEngine.Close()
|
||||
|
||||
st, err := store.New(cfg.RedisURL, cfg.JobResultTTL)
|
||||
if err != nil {
|
||||
log.Fatalf("redis: %v", err)
|
||||
}
|
||||
defer st.Close()
|
||||
|
||||
stg, err := storage.New(cfg.MinIOEndpoint, cfg.MinIOAccessKey, cfg.MinIOSecretKey, cfg.MinIOBucket, cfg.MinIOUseSSL)
|
||||
if err != nil {
|
||||
log.Fatalf("minio: %v", err)
|
||||
}
|
||||
initCtx, initCancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
if err := stg.EnsureBucket(initCtx); err != nil {
|
||||
initCancel()
|
||||
log.Fatalf("minio bucket: %v", err)
|
||||
}
|
||||
initCancel()
|
||||
|
||||
pub, err := queue.NewPublisher(cfg.NATSURL, cfg.StreamName, cfg.SubjectPrefix)
|
||||
if err != nil {
|
||||
log.Fatalf("nats: %v", err)
|
||||
}
|
||||
|
||||
openAPIPath := os.Getenv("OPENAPI_PATH")
|
||||
if openAPIPath == "" {
|
||||
openAPIPath = "api/openapi.json"
|
||||
openAPIPath = "api/openapi.yaml"
|
||||
}
|
||||
|
||||
srv := api.NewServer(cfg, predictor, ocrEngine, st, stg, pub, api.LoadOpenAPI(openAPIPath))
|
||||
srv := api.NewServer(cfg, predictor, ocrEngine, api.LoadOpenAPI(openAPIPath))
|
||||
|
||||
addr := ":" + cfg.Port
|
||||
log.Printf("api listening on %s", addr)
|
||||
log.Printf("count-api escuchando en %s", addr)
|
||||
if err := api.ListenAndServe(ctx, addr, srv.Router()); err != nil {
|
||||
log.Fatalf("api: %v", err)
|
||||
log.Fatalf("server: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,157 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"math"
|
||||
"os"
|
||||
"os/exec"
|
||||
"os/signal"
|
||||
"strconv"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/pys/pdf-count/internal/store"
|
||||
)
|
||||
|
||||
func main() {
|
||||
cfg := loadConfig()
|
||||
ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
||||
defer cancel()
|
||||
|
||||
st, err := store.New(cfg.redisURL, 72*time.Hour)
|
||||
if err != nil {
|
||||
log.Fatalf("redis: %v", err)
|
||||
}
|
||||
defer st.Close()
|
||||
|
||||
log.Printf("autoscaler started min=%d max=%d slots=%d service=%s",
|
||||
cfg.minWorkers, cfg.maxWorkers, cfg.workerSlots, cfg.composeService)
|
||||
|
||||
var idleSince time.Time
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
log.Println("autoscaler stopped")
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
pending, err := st.PendingCount(ctx)
|
||||
if err != nil {
|
||||
log.Printf("autoscaler pending count: %v", err)
|
||||
time.Sleep(cfg.pollInterval)
|
||||
continue
|
||||
}
|
||||
|
||||
current, err := countWorkers(cfg)
|
||||
if err != nil {
|
||||
log.Printf("autoscaler count workers: %v", err)
|
||||
time.Sleep(cfg.pollInterval)
|
||||
continue
|
||||
}
|
||||
|
||||
desired := cfg.minWorkers
|
||||
if pending > 0 {
|
||||
desired = int(math.Ceil(float64(pending) / float64(cfg.workerSlots)))
|
||||
if desired < cfg.minWorkers {
|
||||
desired = cfg.minWorkers
|
||||
}
|
||||
if desired > cfg.maxWorkers {
|
||||
desired = cfg.maxWorkers
|
||||
}
|
||||
idleSince = time.Time{}
|
||||
} else if current > cfg.minWorkers {
|
||||
if idleSince.IsZero() {
|
||||
idleSince = time.Now()
|
||||
} else if time.Since(idleSince) >= cfg.scaleDownIdle {
|
||||
desired = current - 1
|
||||
idleSince = time.Time{}
|
||||
} else {
|
||||
desired = current
|
||||
}
|
||||
}
|
||||
|
||||
if desired != current {
|
||||
log.Printf("autoscaler pending=%d current=%d desired=%d scaling", pending, current, desired)
|
||||
if err := scaleWorkers(cfg, desired); err != nil {
|
||||
log.Printf("autoscaler scale: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
time.Sleep(cfg.pollInterval)
|
||||
}
|
||||
}
|
||||
|
||||
type autoscalerConfig struct {
|
||||
redisURL string
|
||||
composeFile string
|
||||
composeService string
|
||||
minWorkers int
|
||||
maxWorkers int
|
||||
workerSlots int
|
||||
pollInterval time.Duration
|
||||
scaleDownIdle time.Duration
|
||||
}
|
||||
|
||||
func loadConfig() autoscalerConfig {
|
||||
return autoscalerConfig{
|
||||
redisURL: getenv("REDIS_URL", "redis://redis:6379/0"),
|
||||
composeFile: getenv("COMPOSE_FILE", "/project/docker-compose.yml"),
|
||||
composeService: getenv("COMPOSE_SERVICE", "worker"),
|
||||
minWorkers: getenvInt("MIN_WORKERS", 1),
|
||||
maxWorkers: getenvInt("MAX_WORKERS", 8),
|
||||
workerSlots: getenvInt("WORKER_SLOTS", 4),
|
||||
pollInterval: time.Duration(getenvInt("POLL_INTERVAL_SEC", 5)) * time.Second,
|
||||
scaleDownIdle: time.Duration(getenvInt("SCALE_DOWN_IDLE_SEC", 60)) * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
func countWorkers(cfg autoscalerConfig) (int, error) {
|
||||
cmd := exec.Command("docker", "compose",
|
||||
"-f", cfg.composeFile,
|
||||
"ps", cfg.composeService,
|
||||
"--status", "running",
|
||||
"-q",
|
||||
)
|
||||
out, err := cmd.Output()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
trimmed := strings.TrimSpace(string(out))
|
||||
if trimmed == "" {
|
||||
return 0, nil
|
||||
}
|
||||
return len(strings.Split(trimmed, "\n")), nil
|
||||
}
|
||||
|
||||
func scaleWorkers(cfg autoscalerConfig, n int) error {
|
||||
scaleArg := cfg.composeService + "=" + strconv.Itoa(n)
|
||||
cmd := exec.Command("docker", "compose",
|
||||
"-f", cfg.composeFile,
|
||||
"up", "-d", "--no-recreate",
|
||||
"--scale", scaleArg,
|
||||
cfg.composeService,
|
||||
)
|
||||
cmd.Stdout = os.Stdout
|
||||
cmd.Stderr = os.Stderr
|
||||
return cmd.Run()
|
||||
}
|
||||
|
||||
func getenv(key, fallback string) string {
|
||||
if v := os.Getenv(key); v != "" {
|
||||
return v
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func getenvInt(key string, fallback int) int {
|
||||
if v := os.Getenv(key); v != "" {
|
||||
if n, err := strconv.Atoi(v); err == nil {
|
||||
return n
|
||||
}
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
@@ -1,7 +0,0 @@
|
||||
package main
|
||||
|
||||
import "github.com/pys/pdf-count/internal/worker"
|
||||
|
||||
func main() {
|
||||
worker.RunMain()
|
||||
}
|
||||
@@ -1,40 +1,20 @@
|
||||
module github.com/pys/pdf-count
|
||||
|
||||
go 1.23.0
|
||||
|
||||
toolchain go1.23.4
|
||||
go 1.25.0
|
||||
|
||||
require (
|
||||
github.com/gen2brain/go-fitz v1.24.14
|
||||
github.com/go-chi/chi/v5 v5.2.1
|
||||
github.com/google/uuid v1.6.0
|
||||
github.com/minio/minio-go/v7 v7.0.90
|
||||
github.com/nats-io/nats.go v1.39.1
|
||||
github.com/otiai10/gosseract/v2 v2.4.1
|
||||
github.com/redis/go-redis/v9 v9.7.3
|
||||
github.com/yalue/onnxruntime_go v1.21.0
|
||||
gopkg.in/yaml.v3 v3.0.1
|
||||
)
|
||||
|
||||
require github.com/rogpeppe/go-internal v1.15.0 // indirect
|
||||
|
||||
require (
|
||||
github.com/cespare/xxhash/v2 v2.2.0 // indirect
|
||||
github.com/davecgh/go-spew v1.1.1 // indirect
|
||||
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
|
||||
github.com/dustin/go-humanize v1.0.1 // indirect
|
||||
github.com/ebitengine/purego v0.8.0 // indirect
|
||||
github.com/go-ini/ini v1.67.0 // indirect
|
||||
github.com/goccy/go-json v0.10.5 // indirect
|
||||
github.com/jupiterrider/ffi v0.2.0 // indirect
|
||||
github.com/klauspost/compress v1.18.0 // indirect
|
||||
github.com/klauspost/cpuid/v2 v2.2.10 // indirect
|
||||
github.com/minio/crc64nvme v1.0.1 // indirect
|
||||
github.com/minio/md5-simd v1.1.2 // indirect
|
||||
github.com/nats-io/nkeys v0.4.9 // indirect
|
||||
github.com/nats-io/nuid v1.0.1 // indirect
|
||||
github.com/pmezard/go-difflib v1.0.0 // indirect
|
||||
github.com/rs/xid v1.6.0 // indirect
|
||||
golang.org/x/crypto v0.36.0 // indirect
|
||||
golang.org/x/net v0.38.0 // indirect
|
||||
golang.org/x/sys v0.31.0 // indirect
|
||||
golang.org/x/text v0.23.0 // indirect
|
||||
github.com/kr/pretty v0.3.0 // indirect
|
||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
|
||||
)
|
||||
|
||||
@@ -1,69 +1,33 @@
|
||||
github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
|
||||
github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c=
|
||||
github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
|
||||
github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0=
|
||||
github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
|
||||
github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
|
||||
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
|
||||
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
|
||||
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
|
||||
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
|
||||
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
|
||||
github.com/ebitengine/purego v0.8.0 h1:JbqvnEzRvPpxhCJzJJ2y0RbiZ8nyjccVUrSM3q+GvvE=
|
||||
github.com/ebitengine/purego v0.8.0/go.mod h1:iIjxzd6CiRiOG0UyXP+V1+jWqUXVjPKLAI0mRfJZTmQ=
|
||||
github.com/gen2brain/go-fitz v1.24.14 h1:09weRkjVtLYNGo7l0J7DyOwBExbwi8SJ9h8YPhw9WEo=
|
||||
github.com/gen2brain/go-fitz v1.24.14/go.mod h1:0KaZeQgASc20Yp5R/pFzyy7SmP01XcoHKNF842U2/S4=
|
||||
github.com/go-chi/chi/v5 v5.2.1 h1:KOIHODQj58PmL80G2Eak4WdvUzjSJSm0vG72crDCqb8=
|
||||
github.com/go-chi/chi/v5 v5.2.1/go.mod h1:L2yAIGWB3H+phAw1NxKwWM+7eUH/lU8pOMm5hHcoops=
|
||||
github.com/go-ini/ini v1.67.0 h1:z6ZrTEZqSWOTyH2FlglNbNgARyHG8oLW9gMELqKr06A=
|
||||
github.com/go-ini/ini v1.67.0/go.mod h1:ByCAeIL28uOIIG0E3PJtZPDL8WnHpFKFOtgjp+3Ies8=
|
||||
github.com/goccy/go-json v0.10.5 h1:Fq85nIqj+gXn/S5ahsiTlK3TmC85qgirsdTP/+DeaC4=
|
||||
github.com/goccy/go-json v0.10.5/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/jupiterrider/ffi v0.2.0 h1:tMM70PexgYNmV+WyaYhJgCvQAvtTCs3wXeILPutihnA=
|
||||
github.com/jupiterrider/ffi v0.2.0/go.mod h1:yqYqX5DdEccAsHeMn+6owkoI2llBLySVAF8dwCDZPVs=
|
||||
github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo=
|
||||
github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ=
|
||||
github.com/klauspost/cpuid/v2 v2.0.1/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
|
||||
github.com/klauspost/cpuid/v2 v2.2.10 h1:tBs3QSyvjDyFTq3uoc/9xFpCuOsJQFNPiAhYdw2skhE=
|
||||
github.com/klauspost/cpuid/v2 v2.2.10/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0=
|
||||
github.com/minio/crc64nvme v1.0.1 h1:DHQPrYPdqK7jQG/Ls5CTBZWeex/2FMS3G5XGkycuFrY=
|
||||
github.com/minio/crc64nvme v1.0.1/go.mod h1:eVfm2fAzLlxMdUGc0EEBGSMmPwmXD5XiNRpnu9J3bvg=
|
||||
github.com/minio/md5-simd v1.1.2 h1:Gdi1DZK69+ZVMoNHRXJyNcxrMA4dSxoYHZSQbirFg34=
|
||||
github.com/minio/md5-simd v1.1.2/go.mod h1:MzdKDxYpY2BT9XQFocsiZf/NKVtR7nkE4RoEpN+20RM=
|
||||
github.com/minio/minio-go/v7 v7.0.90 h1:TmSj1083wtAD0kEYTx7a5pFsv3iRYMsOJ6A4crjA1lE=
|
||||
github.com/minio/minio-go/v7 v7.0.90/go.mod h1:uvMUcGrpgeSAAI6+sD3818508nUyMULw94j2Nxku/Go=
|
||||
github.com/nats-io/nats.go v1.39.1 h1:oTkfKBmz7W047vRxV762M67ZdXeOtUgvbBaNoQ+3PPk=
|
||||
github.com/nats-io/nats.go v1.39.1/go.mod h1:MgRb8oOdigA6cYpEPhXJuRVH6UE/V4jblJ2jQ27IXYM=
|
||||
github.com/nats-io/nkeys v0.4.9 h1:qe9Faq2Gxwi6RZnZMXfmGMZkg3afLLOtrU+gDZJ35b0=
|
||||
github.com/nats-io/nkeys v0.4.9/go.mod h1:jcMqs+FLG+W5YO36OX6wFIFcmpdAns+w1Wm6D3I/evE=
|
||||
github.com/nats-io/nuid v1.0.1 h1:5iA8DT8V7q8WK2EScv2padNa/rTESc1KdnPw4TC2paw=
|
||||
github.com/nats-io/nuid v1.0.1/go.mod h1:19wcPz3Ph3q0Jbyiqsd0kePYG7A95tJPxeL+1OSON2c=
|
||||
github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
|
||||
github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
|
||||
github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0=
|
||||
github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk=
|
||||
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
|
||||
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
|
||||
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
|
||||
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
|
||||
github.com/otiai10/gosseract/v2 v2.4.1 h1:G8AyBpXEeSlcq8TI85LH/pM5SXk8Djy2GEXisgyblRw=
|
||||
github.com/otiai10/gosseract/v2 v2.4.1/go.mod h1:1gNWP4Hgr2o7yqWfs6r5bZxAatjOIdqWxJLWsTsembk=
|
||||
github.com/otiai10/mint v1.6.3 h1:87qsV/aw1F5as1eH1zS/yqHY85ANKVMgkDrf9rcxbQs=
|
||||
github.com/otiai10/mint v1.6.3/go.mod h1:MJm72SBthJjz8qhefc4z1PYEieWmy8Bku7CjcAqyUSM=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/redis/go-redis/v9 v9.7.3 h1:YpPyAayJV+XErNsatSElgRZZVCwXX9QzkKYNvO7x0wM=
|
||||
github.com/redis/go-redis/v9 v9.7.3/go.mod h1:bGUrSggJ9X9GUmZpZNEOQKaANxSGgOEBRltRTZHSvrA=
|
||||
github.com/rs/xid v1.6.0 h1:fV591PaemRlL6JfRxGDEPl69wICngIQ3shQtzfy2gxU=
|
||||
github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0=
|
||||
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
|
||||
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
|
||||
github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc=
|
||||
github.com/rogpeppe/go-internal v1.15.0 h1:D0RCU5rMAp+SpgkiNdrjfJ+LX4J1M32V2NeCY7EJ6hc=
|
||||
github.com/rogpeppe/go-internal v1.15.0/go.mod h1:DrUVZyrJU+txYW5/1kwtXQSMFio52ZOxX7yM1VHvnxs=
|
||||
github.com/yalue/onnxruntime_go v1.21.0 h1:DdtvfY7OP5gR8mwPDqAOAQckf+KcI30hPNJL8hQaYWI=
|
||||
github.com/yalue/onnxruntime_go v1.21.0/go.mod h1:b4X26A8pekNb1ACJ58wAXgNKeUCGEAQ9dmACut9Sm/4=
|
||||
golang.org/x/crypto v0.36.0 h1:AnAEvhDddvBdpY+uR+MyHmuZzzNqXSe/GvuDeob5L34=
|
||||
golang.org/x/crypto v0.36.0/go.mod h1:Y4J0ReaxCR1IMaabaSMugxJES1EpwhBHhv2bDHklZvc=
|
||||
golang.org/x/net v0.38.0 h1:vRMAPTMaeGqVhG5QyLJHqNDwecKTomGeqbnfZyKlBI8=
|
||||
golang.org/x/net v0.38.0/go.mod h1:ivrbrMbzFq5J41QOQh0siUuly180yBYtLp+CKbEaFx8=
|
||||
golang.org/x/sys v0.31.0 h1:ioabZlmFYtWhL+TRYpcnNlLwhyxaM9kWTDEmfnprqik=
|
||||
golang.org/x/sys v0.31.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
|
||||
golang.org/x/text v0.23.0 h1:D71I7dUrlY+VX0gQShAThNGHFxZ13dGLBHQLVl1mJlY=
|
||||
golang.org/x/text v0.23.0/go.mod h1:/BLNzu4aZCJ1+kcD0DNRotWKage4q2rGVAg4o22unh4=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
|
||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
|
||||
gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
|
||||
@@ -1,53 +0,0 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/pys/pdf-count/internal/pipeline"
|
||||
"github.com/pys/pdf-count/internal/store"
|
||||
)
|
||||
|
||||
// FilePayload holds a PDF ready for inline job processing.
|
||||
type FilePayload struct {
|
||||
Filename string
|
||||
Data []byte
|
||||
}
|
||||
|
||||
// JobRunner processes uploaded job files in-process using the API OCR pool.
|
||||
type JobRunner struct {
|
||||
store *store.Store
|
||||
newProc func(includeRejected bool) *pipeline.Processor
|
||||
sem chan struct{}
|
||||
}
|
||||
|
||||
func NewJobRunner(st *store.Store, newProc func(includeRejected bool) *pipeline.Processor, concurrency int) *JobRunner {
|
||||
if concurrency < 1 {
|
||||
concurrency = 1
|
||||
}
|
||||
return &JobRunner{
|
||||
store: st,
|
||||
newProc: newProc,
|
||||
sem: make(chan struct{}, concurrency),
|
||||
}
|
||||
}
|
||||
|
||||
// ProcessOne schedules a single PDF for inline OCR (non-blocking).
|
||||
func (jr *JobRunner) ProcessOne(jobID string, file FilePayload, includeRejected bool) {
|
||||
go func() {
|
||||
jr.sem <- struct{}{}
|
||||
defer func() { <-jr.sem }()
|
||||
|
||||
ctx := context.Background()
|
||||
proc := jr.newProc(includeRejected)
|
||||
result := proc.ProcessPDF(file.Filename, file.Data)
|
||||
_ = jr.store.SaveResult(ctx, jobID, result)
|
||||
_ = jr.store.TryCompleteJob(ctx, jobID)
|
||||
}()
|
||||
}
|
||||
|
||||
// ProcessFiles runs OCR in background goroutines and updates job progress in Redis.
|
||||
func (jr *JobRunner) ProcessFiles(jobID string, files []FilePayload, includeRejected bool) {
|
||||
for _, file := range files {
|
||||
jr.ProcessOne(jobID, file, includeRejected)
|
||||
}
|
||||
}
|
||||
@@ -4,87 +4,47 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
"github.com/go-chi/chi/v5/middleware"
|
||||
"github.com/google/uuid"
|
||||
"github.com/pys/pdf-count/internal/config"
|
||||
"github.com/pys/pdf-count/internal/domain"
|
||||
"github.com/pys/pdf-count/internal/model"
|
||||
"github.com/pys/pdf-count/internal/ocr"
|
||||
"github.com/pys/pdf-count/internal/pipeline"
|
||||
"github.com/pys/pdf-count/internal/queue"
|
||||
"github.com/pys/pdf-count/internal/storage"
|
||||
"github.com/pys/pdf-count/internal/store"
|
||||
"github.com/pys/pdf-count/web"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
// Server holds the shared, stateless dependencies of the count-api.
|
||||
type Server struct {
|
||||
cfg config.Config
|
||||
predictor *model.Predictor
|
||||
predictor *model.Predictor // nil → heuristic fallback
|
||||
ocrEngine *ocr.Engine
|
||||
store *store.Store
|
||||
storage *storage.Client
|
||||
publisher *queue.Publisher
|
||||
jobRunner *JobRunner
|
||||
openAPI []byte
|
||||
ready bool
|
||||
readyMu sync.RWMutex
|
||||
}
|
||||
|
||||
func NewServer(cfg config.Config, predictor *model.Predictor, ocrEngine *ocr.Engine, st *store.Store, stg *storage.Client, pub *queue.Publisher, openAPI []byte) *Server {
|
||||
srv := &Server{
|
||||
// NewServer constructs a Server. predictor may be nil when the ONNX model is
|
||||
// unavailable; the pipeline will fall back to a counting heuristic.
|
||||
func NewServer(
|
||||
cfg config.Config,
|
||||
predictor *model.Predictor,
|
||||
ocrEngine *ocr.Engine,
|
||||
openAPI []byte,
|
||||
) *Server {
|
||||
return &Server{
|
||||
cfg: cfg,
|
||||
predictor: predictor,
|
||||
ocrEngine: ocrEngine,
|
||||
store: st,
|
||||
storage: stg,
|
||||
publisher: pub,
|
||||
openAPI: openAPI,
|
||||
ready: true,
|
||||
}
|
||||
srv.jobRunner = NewJobRunner(st, srv.newProcessor, srv.ocrConcurrency())
|
||||
return srv
|
||||
}
|
||||
|
||||
func (s *Server) processConcurrency() int {
|
||||
c := s.cfg.SyncProcessConcurrency
|
||||
if c <= 0 {
|
||||
c = s.cfg.OCRPoolSize
|
||||
}
|
||||
if c <= 0 {
|
||||
c = 4
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
func (s *Server) ocrConcurrency() int {
|
||||
c := s.processConcurrency()
|
||||
pool := s.cfg.OCRPoolSize
|
||||
if pool <= 0 {
|
||||
pool = 4
|
||||
}
|
||||
if c > pool {
|
||||
c = pool
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
func (s *Server) SetReady(v bool) {
|
||||
s.readyMu.Lock()
|
||||
s.ready = v
|
||||
s.readyMu.Unlock()
|
||||
}
|
||||
|
||||
// Router registers all HTTP routes.
|
||||
func (s *Server) Router() http.Handler {
|
||||
r := chi.NewRouter()
|
||||
r.Use(middleware.RequestID)
|
||||
@@ -92,46 +52,41 @@ func (s *Server) Router() http.Handler {
|
||||
r.Use(middleware.Recoverer)
|
||||
|
||||
r.Get("/health", s.health)
|
||||
r.Get("/ready", s.readyHandler)
|
||||
r.Get("/ready", s.ready)
|
||||
r.Get("/openapi.json", s.openAPIHandler)
|
||||
|
||||
r.Get("/demo", func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Redirect(w, r, "/demo/", http.StatusMovedPermanently)
|
||||
})
|
||||
r.Handle("/demo/*", http.StripPrefix("/demo/", http.FileServer(http.FS(web.FS))))
|
||||
|
||||
r.Route("/v1", func(r chi.Router) {
|
||||
r.Post("/count", s.countSync)
|
||||
r.Post("/jobs", s.createJob)
|
||||
r.Post("/jobs/{jobID}/upload", s.uploadChunk)
|
||||
r.Get("/jobs/{jobID}", s.getJob)
|
||||
r.Get("/jobs/{jobID}/results", s.getJobResults)
|
||||
r.Delete("/jobs/{jobID}", s.cancelJob)
|
||||
})
|
||||
|
||||
return r
|
||||
}
|
||||
|
||||
func (s *Server) health(w http.ResponseWriter, r *http.Request) {
|
||||
// health returns a simple liveness response.
|
||||
func (s *Server) health(w http.ResponseWriter, _ *http.Request) {
|
||||
writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
|
||||
}
|
||||
|
||||
func (s *Server) readyHandler(w http.ResponseWriter, r *http.Request) {
|
||||
s.readyMu.RLock()
|
||||
ok := s.ready
|
||||
s.readyMu.RUnlock()
|
||||
if !ok {
|
||||
writeJSON(w, http.StatusServiceUnavailable, map[string]string{"status": "not_ready"})
|
||||
return
|
||||
}
|
||||
// ready returns 200 once the server is accepting traffic (always ready after init).
|
||||
func (s *Server) ready(w http.ResponseWriter, _ *http.Request) {
|
||||
writeJSON(w, http.StatusOK, map[string]string{"status": "ready"})
|
||||
}
|
||||
|
||||
func (s *Server) openAPIHandler(w http.ResponseWriter, r *http.Request) {
|
||||
// openAPIHandler serves the OpenAPI specification as JSON.
|
||||
func (s *Server) openAPIHandler(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write(s.openAPI)
|
||||
}
|
||||
|
||||
// countSync processes a multipart batch of PDFs synchronously and returns the
|
||||
// counting results in a single JSON response. No state is persisted.
|
||||
//
|
||||
// POST /v1/count
|
||||
//
|
||||
// Query: include_rejected=true|false
|
||||
// Body: multipart/form-data, field "files" (one or more PDFs)
|
||||
// Limits: SYNC_MAX_FILES (default 50), MAX_PDF_SIZE_MB (default 20)
|
||||
func (s *Server) countSync(w http.ResponseWriter, r *http.Request) {
|
||||
includeRejected := r.URL.Query().Get("include_rejected") == "true"
|
||||
maxBytes := int64(s.cfg.MaxPDFSizeMB) * 1024 * 1024
|
||||
@@ -151,61 +106,51 @@ func (s *Server) countSync(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
proc := s.newProcessor(includeRejected)
|
||||
type fileInput struct {
|
||||
filename string
|
||||
data []byte
|
||||
errRes *domain.PDFResult
|
||||
}
|
||||
inputs := make([]fileInput, 0, len(files))
|
||||
|
||||
inputs := make([]fileInput, 0, len(files))
|
||||
for _, fh := range files {
|
||||
name := filepath.Base(fh.Filename)
|
||||
f, err := fh.Open()
|
||||
if err != nil {
|
||||
errRes := domain.PDFResult{
|
||||
Filename: filepath.Base(fh.Filename),
|
||||
Status: domain.StatusError,
|
||||
Error: err.Error(),
|
||||
}
|
||||
errRes := domain.PDFResult{Filename: name, Status: domain.StatusError, Error: err.Error()}
|
||||
inputs = append(inputs, fileInput{errRes: &errRes})
|
||||
continue
|
||||
}
|
||||
data, err := io.ReadAll(io.LimitReader(f, maxBytes+1))
|
||||
_ = f.Close()
|
||||
if err != nil {
|
||||
errRes := domain.PDFResult{
|
||||
Filename: filepath.Base(fh.Filename),
|
||||
Status: domain.StatusError,
|
||||
Error: err.Error(),
|
||||
}
|
||||
errRes := domain.PDFResult{Filename: name, Status: domain.StatusError, Error: err.Error()}
|
||||
inputs = append(inputs, fileInput{errRes: &errRes})
|
||||
continue
|
||||
}
|
||||
if int64(len(data)) > maxBytes {
|
||||
errRes := domain.PDFResult{
|
||||
Filename: filepath.Base(fh.Filename),
|
||||
Status: domain.StatusError,
|
||||
Error: "file too large",
|
||||
}
|
||||
errRes := domain.PDFResult{Filename: name, Status: domain.StatusError, Error: "file too large"}
|
||||
inputs = append(inputs, fileInput{errRes: &errRes})
|
||||
continue
|
||||
}
|
||||
inputs = append(inputs, fileInput{
|
||||
filename: filepath.Base(fh.Filename),
|
||||
data: data,
|
||||
})
|
||||
inputs = append(inputs, fileInput{filename: name, data: data})
|
||||
}
|
||||
|
||||
proc := s.newProcessor(includeRejected)
|
||||
results := make([]domain.PDFResult, len(inputs))
|
||||
failed := 0
|
||||
sem := make(chan struct{}, s.ocrConcurrency())
|
||||
|
||||
concurrency := s.syncConcurrency()
|
||||
sem := make(chan struct{}, concurrency)
|
||||
var wg sync.WaitGroup
|
||||
var mu sync.Mutex
|
||||
|
||||
for i, in := range inputs {
|
||||
if in.errRes != nil {
|
||||
results[i] = *in.errRes
|
||||
mu.Lock()
|
||||
failed++
|
||||
mu.Unlock()
|
||||
continue
|
||||
}
|
||||
wg.Add(1)
|
||||
@@ -231,314 +176,23 @@ func (s *Server) countSync(w http.ResponseWriter, r *http.Request) {
|
||||
})
|
||||
}
|
||||
|
||||
type createJobRequest struct {
|
||||
Files []string `json:"files"`
|
||||
Prefix string `json:"prefix"`
|
||||
IncludeRejected bool `json:"include_rejected"`
|
||||
CallbackURL string `json:"callback_url"`
|
||||
}
|
||||
|
||||
func (s *Server) createJob(w http.ResponseWriter, r *http.Request) {
|
||||
var req createJobRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
httpError(w, http.StatusBadRequest, "invalid json body")
|
||||
return
|
||||
// syncConcurrency returns the effective parallelism for /v1/count requests.
|
||||
func (s *Server) syncConcurrency() int {
|
||||
c := s.cfg.SyncProcessConcurrency
|
||||
if c <= 0 {
|
||||
c = s.cfg.OCRPoolSize
|
||||
}
|
||||
|
||||
ctx := r.Context()
|
||||
pending, err := s.store.PendingCount(ctx)
|
||||
if err == nil && int(pending) > s.cfg.QueueMaxPending {
|
||||
httpError(w, http.StatusTooManyRequests, "queue full")
|
||||
return
|
||||
if c <= 0 {
|
||||
c = 4
|
||||
}
|
||||
|
||||
jobID := uuid.NewString()
|
||||
var objectKeys []string
|
||||
var filenames []string
|
||||
|
||||
if req.Prefix != "" {
|
||||
prefix := strings.TrimSuffix(req.Prefix, "/")
|
||||
keys, err := s.storage.ListByPrefix(ctx, prefix)
|
||||
if err != nil {
|
||||
httpError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
objectKeys = keys
|
||||
for _, k := range keys {
|
||||
filenames = append(filenames, filepath.Base(k))
|
||||
}
|
||||
} else if len(req.Files) > 0 {
|
||||
for _, f := range req.Files {
|
||||
name := filepath.Base(f)
|
||||
objectKeys = append(objectKeys, storage.JobObjectKey(jobID, name))
|
||||
filenames = append(filenames, name)
|
||||
}
|
||||
} else {
|
||||
// Empty job shell for chunked upload flow.
|
||||
if err := s.store.CreateJob(ctx, jobID, 0, req.IncludeRejected); err != nil {
|
||||
httpError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusAccepted, domain.JobStatus{
|
||||
JobID: jobID,
|
||||
Status: domain.JobQueued,
|
||||
Total: 0,
|
||||
})
|
||||
return
|
||||
pool := s.cfg.OCRPoolSize
|
||||
if pool <= 0 {
|
||||
pool = 4
|
||||
}
|
||||
|
||||
if len(objectKeys) > s.cfg.JobMaxFiles {
|
||||
httpError(w, http.StatusBadRequest, "too many files in job")
|
||||
return
|
||||
if c > pool {
|
||||
c = pool
|
||||
}
|
||||
|
||||
if err := s.store.CreateJob(ctx, jobID, len(objectKeys), req.IncludeRejected); err != nil {
|
||||
httpError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
if err := s.store.AddJobFiles(ctx, jobID, filenames); err != nil {
|
||||
httpError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
_ = s.store.MarkRunning(ctx, jobID)
|
||||
|
||||
for i, key := range objectKeys {
|
||||
msg := domain.QueueMessage{
|
||||
JobID: jobID,
|
||||
ObjectKey: key,
|
||||
Filename: filenames[i],
|
||||
IncludeRejected: req.IncludeRejected,
|
||||
}
|
||||
if err := s.publisher.Publish(ctx, msg); err != nil {
|
||||
httpError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
writeJSON(w, http.StatusAccepted, domain.JobStatus{
|
||||
JobID: jobID,
|
||||
Status: domain.JobRunning,
|
||||
Total: len(objectKeys),
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) uploadChunk(w http.ResponseWriter, r *http.Request) {
|
||||
jobID := chi.URLParam(r, "jobID")
|
||||
includeRejected := r.URL.Query().Get("include_rejected") == "true"
|
||||
maxBytes := int64(s.cfg.MaxPDFSizeMB) * 1024 * 1024
|
||||
|
||||
if err := r.ParseMultipartForm(maxBytes * 500); err != nil {
|
||||
httpError(w, http.StatusBadRequest, "invalid multipart form")
|
||||
return
|
||||
}
|
||||
files := r.MultipartForm.File["files"]
|
||||
if len(files) == 0 {
|
||||
httpError(w, http.StatusBadRequest, "field 'files' required")
|
||||
return
|
||||
}
|
||||
if len(files) > 500 {
|
||||
httpError(w, http.StatusBadRequest, "max 500 files per chunk")
|
||||
return
|
||||
}
|
||||
|
||||
ctx := r.Context()
|
||||
status, err := s.store.GetJob(ctx, jobID)
|
||||
if err != nil {
|
||||
_ = s.store.CreateJob(ctx, jobID, 0, includeRejected)
|
||||
status = domain.JobStatus{JobID: jobID, Status: domain.JobQueued}
|
||||
}
|
||||
if status.Status == domain.JobCancelled {
|
||||
httpError(w, http.StatusConflict, "job cancelled")
|
||||
return
|
||||
}
|
||||
|
||||
var filenames []string
|
||||
|
||||
if s.cfg.JobInlineProcessing {
|
||||
_ = s.store.MarkRunning(ctx, jobID)
|
||||
for _, fh := range files {
|
||||
f, err := fh.Open()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
data, err := io.ReadAll(io.LimitReader(f, maxBytes+1))
|
||||
_ = f.Close()
|
||||
if err != nil || int64(len(data)) > maxBytes {
|
||||
continue
|
||||
}
|
||||
name := filepath.Base(fh.Filename)
|
||||
if err := s.store.AddJobFiles(ctx, jobID, []string{name}); err != nil {
|
||||
httpError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
if err := s.store.IncrementTotal(ctx, jobID, 1); err != nil {
|
||||
httpError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
filenames = append(filenames, name)
|
||||
s.jobRunner.ProcessOne(jobID, FilePayload{Filename: name, Data: data}, includeRejected)
|
||||
}
|
||||
} else {
|
||||
var err error
|
||||
filenames, err = s.uploadChunkMinIO(ctx, jobID, files, includeRejected, maxBytes)
|
||||
if err != nil {
|
||||
httpError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if len(filenames) == 0 {
|
||||
httpError(w, http.StatusBadRequest, "no valid files uploaded")
|
||||
return
|
||||
}
|
||||
|
||||
writeJSON(w, http.StatusAccepted, map[string]interface{}{
|
||||
"job_id": jobID,
|
||||
"uploaded": len(filenames),
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) uploadChunkMinIO(ctx context.Context, jobID string, files []*multipart.FileHeader, includeRejected bool, maxBytes int64) ([]string, error) {
|
||||
filenames := make([]string, 0, len(files))
|
||||
for _, fh := range files {
|
||||
filenames = append(filenames, filepath.Base(fh.Filename))
|
||||
}
|
||||
if len(filenames) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
if err := s.store.AddJobFiles(ctx, jobID, filenames); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.store.IncrementTotal(ctx, jobID, len(filenames)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.store.MarkRunning(ctx, jobID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
ioLimit := 8
|
||||
if len(files) < ioLimit {
|
||||
ioLimit = len(files)
|
||||
}
|
||||
ioSem := make(chan struct{}, ioLimit)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
var mu sync.Mutex
|
||||
uploadedNames := make([]string, 0, len(files))
|
||||
|
||||
for _, fh := range files {
|
||||
wg.Add(1)
|
||||
go func(fh *multipart.FileHeader) {
|
||||
defer wg.Done()
|
||||
ioSem <- struct{}{}
|
||||
defer func() { <-ioSem }()
|
||||
|
||||
name := filepath.Base(fh.Filename)
|
||||
f, err := fh.Open()
|
||||
if err != nil {
|
||||
s.recordUploadError(ctx, jobID, name, err.Error())
|
||||
return
|
||||
}
|
||||
data, err := io.ReadAll(io.LimitReader(f, maxBytes+1))
|
||||
_ = f.Close()
|
||||
if err != nil {
|
||||
s.recordUploadError(ctx, jobID, name, err.Error())
|
||||
return
|
||||
}
|
||||
if int64(len(data)) > maxBytes {
|
||||
s.recordUploadError(ctx, jobID, name, "file too large")
|
||||
return
|
||||
}
|
||||
|
||||
key := storage.JobObjectKey(jobID, name)
|
||||
if err := s.storage.Put(ctx, key, data, "application/pdf"); err != nil {
|
||||
s.recordUploadError(ctx, jobID, name, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
msg := domain.QueueMessage{
|
||||
JobID: jobID,
|
||||
ObjectKey: key,
|
||||
Filename: name,
|
||||
IncludeRejected: includeRejected,
|
||||
}
|
||||
if err := s.publisher.Publish(ctx, msg); err != nil {
|
||||
s.recordUploadError(ctx, jobID, name, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
uploadedNames = append(uploadedNames, name)
|
||||
mu.Unlock()
|
||||
}(fh)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if len(uploadedNames) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
return uploadedNames, nil
|
||||
}
|
||||
|
||||
func (s *Server) recordUploadError(ctx context.Context, jobID, filename, errMsg string) {
|
||||
_ = s.store.SaveResult(ctx, jobID, domain.PDFResult{
|
||||
Filename: filename,
|
||||
Status: domain.StatusError,
|
||||
Error: errMsg,
|
||||
})
|
||||
_ = s.store.TryCompleteJob(ctx, jobID)
|
||||
}
|
||||
|
||||
func (s *Server) getJob(w http.ResponseWriter, r *http.Request) {
|
||||
jobID := chi.URLParam(r, "jobID")
|
||||
waitMs, _ := strconv.Atoi(r.URL.Query().Get("wait_ms"))
|
||||
if waitMs > 30000 {
|
||||
waitMs = 30000
|
||||
}
|
||||
|
||||
deadline := time.Now()
|
||||
if waitMs > 0 {
|
||||
deadline = deadline.Add(time.Duration(waitMs) * time.Millisecond)
|
||||
}
|
||||
|
||||
for {
|
||||
status, err := s.store.GetJob(r.Context(), jobID)
|
||||
if err != nil {
|
||||
httpError(w, http.StatusNotFound, "job not found")
|
||||
return
|
||||
}
|
||||
if isTerminalJobStatus(status.Status) || waitMs <= 0 || !time.Now().Before(deadline) {
|
||||
writeJSON(w, http.StatusOK, status)
|
||||
return
|
||||
}
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
func isTerminalJobStatus(status string) bool {
|
||||
return status == domain.JobCompleted || status == domain.JobCancelled
|
||||
}
|
||||
|
||||
func (s *Server) getJobResults(w http.ResponseWriter, r *http.Request) {
|
||||
jobID := chi.URLParam(r, "jobID")
|
||||
cursor := r.URL.Query().Get("cursor")
|
||||
limit, _ := strconv.Atoi(r.URL.Query().Get("limit"))
|
||||
page, err := s.store.GetResultsPage(r.Context(), jobID, cursor, limit)
|
||||
if err != nil {
|
||||
httpError(w, http.StatusNotFound, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, page)
|
||||
}
|
||||
|
||||
func (s *Server) cancelJob(w http.ResponseWriter, r *http.Request) {
|
||||
jobID := chi.URLParam(r, "jobID")
|
||||
if err := s.store.CancelJob(r.Context(), jobID); err != nil {
|
||||
httpError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]string{"job_id": jobID, "status": domain.JobCancelled})
|
||||
return c
|
||||
}
|
||||
|
||||
func (s *Server) newProcessor(includeRejected bool) *pipeline.Processor {
|
||||
@@ -552,20 +206,24 @@ func (s *Server) newProcessor(includeRejected bool) *pipeline.Processor {
|
||||
})
|
||||
}
|
||||
|
||||
// writeJSON encodes v as JSON and writes it with the given HTTP status.
|
||||
func writeJSON(w http.ResponseWriter, status int, v interface{}) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(status)
|
||||
_ = json.NewEncoder(w).Encode(v)
|
||||
}
|
||||
|
||||
// httpError writes a JSON error response.
|
||||
func httpError(w http.ResponseWriter, status int, msg string) {
|
||||
writeJSON(w, status, map[string]string{"error": msg})
|
||||
}
|
||||
|
||||
// LoadOpenAPI reads an OpenAPI spec from path.
|
||||
// YAML files are converted to JSON on the fly. Returns a minimal stub on error.
|
||||
func LoadOpenAPI(path string) []byte {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return []byte(`{"openapi":"3.1.0","info":{"title":"pdf-count","version":"1.0.0"}}`)
|
||||
return []byte(`{"openapi":"3.1.0","info":{"title":"count-api","version":"1.0.0"}}`)
|
||||
}
|
||||
if strings.HasSuffix(strings.ToLower(path), ".yaml") || strings.HasSuffix(strings.ToLower(path), ".yml") {
|
||||
var doc interface{}
|
||||
@@ -578,6 +236,7 @@ func LoadOpenAPI(path string) []byte {
|
||||
return data
|
||||
}
|
||||
|
||||
// ListenAndServe starts the HTTP server and blocks until ctx is cancelled.
|
||||
func ListenAndServe(ctx context.Context, addr string, handler http.Handler) error {
|
||||
srv := &http.Server{Addr: addr, Handler: handler}
|
||||
go func() {
|
||||
|
||||
@@ -6,67 +6,39 @@ import (
|
||||
"time"
|
||||
)
|
||||
|
||||
// Config holds all runtime parameters for the count-api service.
|
||||
// All values are loaded from environment variables with safe defaults.
|
||||
type Config struct {
|
||||
Port string
|
||||
RedisURL string
|
||||
NATSURL string
|
||||
MinIOEndpoint string
|
||||
MinIOAccessKey string
|
||||
MinIOSecretKey string
|
||||
MinIOBucket string
|
||||
MinIOUseSSL bool
|
||||
ModelPath string
|
||||
InferencePath string
|
||||
OCRDPI int
|
||||
OCRLang string
|
||||
DBSCANEps float64
|
||||
DBSCANMinSamples int
|
||||
SyncMaxFiles int
|
||||
JobMaxFiles int
|
||||
QueueMaxPending int
|
||||
WorkerConcurrency int
|
||||
MaxPDFSizeMB int
|
||||
PDFProcessTimeout time.Duration
|
||||
JobResultTTL time.Duration
|
||||
StreamName string
|
||||
SubjectPrefix string
|
||||
Port string
|
||||
ModelPath string
|
||||
InferencePath string
|
||||
OCRDPI int
|
||||
OCRLang string
|
||||
OCRPoolSize int
|
||||
LogPipelineTiming bool
|
||||
NATSFetchMaxWait time.Duration
|
||||
JobInlineProcessing bool
|
||||
DBSCANEps float64
|
||||
DBSCANMinSamples int
|
||||
SyncMaxFiles int
|
||||
MaxPDFSizeMB int
|
||||
PDFProcessTimeout time.Duration
|
||||
SyncProcessConcurrency int
|
||||
LogPipelineTiming bool
|
||||
}
|
||||
|
||||
func Load() Config {
|
||||
return Config{
|
||||
Port: getenv("PORT", "8080"),
|
||||
RedisURL: getenv("REDIS_URL", "redis://redis:6379/0"),
|
||||
NATSURL: getenv("NATS_URL", "nats://nats:4222"),
|
||||
MinIOEndpoint: getenv("MINIO_ENDPOINT", "minio:9000"),
|
||||
MinIOAccessKey: getenv("MINIO_ACCESS_KEY", "minioadmin"),
|
||||
MinIOSecretKey: getenv("MINIO_SECRET_KEY", "minioadmin"),
|
||||
MinIOBucket: getenv("MINIO_BUCKET", "pdf-count"),
|
||||
MinIOUseSSL: getenvBool("MINIO_USE_SSL", false),
|
||||
ModelPath: getenv("MODEL_PATH", "/app/models/model.onnx"),
|
||||
InferencePath: getenv("INFERENCE_CONFIG", "/app/models/inference.json"),
|
||||
OCRDPI: getenvInt("OCR_DPI", 300),
|
||||
OCRLang: getenv("OCR_LANG", "spa"),
|
||||
DBSCANEps: getenvFloat("DBSCAN_EPS", 10),
|
||||
DBSCANMinSamples: getenvInt("DBSCAN_MIN_SAMPLES", 1),
|
||||
SyncMaxFiles: getenvInt("SYNC_MAX_FILES", 50),
|
||||
JobMaxFiles: getenvInt("JOB_MAX_FILES", 50000),
|
||||
QueueMaxPending: getenvInt("QUEUE_MAX_PENDING", 100000),
|
||||
WorkerConcurrency: getenvInt("WORKER_CONCURRENCY", 4),
|
||||
MaxPDFSizeMB: getenvInt("MAX_PDF_SIZE_MB", 20),
|
||||
PDFProcessTimeout: time.Duration(getenvInt("PDF_PROCESS_TIMEOUT_SEC", 120)) * time.Second,
|
||||
JobResultTTL: time.Duration(getenvInt("JOB_RESULT_TTL_HOURS", 72)) * time.Hour,
|
||||
StreamName: getenv("NATS_STREAM", "PDF_JOBS"),
|
||||
SubjectPrefix: getenv("NATS_SUBJECT_PREFIX", "pdf.process"),
|
||||
Port: getenv("PORT", "8080"),
|
||||
ModelPath: getenv("MODEL_PATH", "/app/models/model.onnx"),
|
||||
InferencePath: getenv("INFERENCE_CONFIG", "/app/models/inference.json"),
|
||||
OCRDPI: getenvInt("OCR_DPI", 300),
|
||||
OCRLang: getenv("OCR_LANG", "spa"),
|
||||
OCRPoolSize: getenvInt("OCR_POOL_SIZE", 0),
|
||||
LogPipelineTiming: getenvBool("LOG_PIPELINE_TIMING", false),
|
||||
NATSFetchMaxWait: time.Duration(getenvInt("NATS_FETCH_MAX_WAIT_MS", 300)) * time.Millisecond,
|
||||
JobInlineProcessing: getenvBool("JOB_INLINE_PROCESSING", false),
|
||||
DBSCANEps: getenvFloat("DBSCAN_EPS", 10),
|
||||
DBSCANMinSamples: getenvInt("DBSCAN_MIN_SAMPLES", 1),
|
||||
SyncMaxFiles: getenvInt("SYNC_MAX_FILES", 50),
|
||||
MaxPDFSizeMB: getenvInt("MAX_PDF_SIZE_MB", 20),
|
||||
PDFProcessTimeout: time.Duration(getenvInt("PDF_PROCESS_TIMEOUT_SEC", 120)) * time.Second,
|
||||
SyncProcessConcurrency: getenvInt("SYNC_PROCESS_CONCURRENCY", 0),
|
||||
LogPipelineTiming: getenvBool("LOG_PIPELINE_TIMING", false),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,13 +1,13 @@
|
||||
package domain
|
||||
|
||||
// ServiceRow is one detected invoice line.
|
||||
// ServiceRow represents a single detected invoice line in a PDF page.
|
||||
type ServiceRow struct {
|
||||
Text string `json:"text"`
|
||||
Probability float64 `json:"probability"`
|
||||
IsService *bool `json:"is_service,omitempty"`
|
||||
}
|
||||
|
||||
// PDFResult is the processing outcome for a single PDF.
|
||||
// PDFResult is the counting outcome for a single PDF document.
|
||||
type PDFResult struct {
|
||||
Filename string `json:"filename"`
|
||||
ServiceCount int `json:"service_count"`
|
||||
@@ -16,44 +16,14 @@ type PDFResult struct {
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// CountResponse wraps sync batch results.
|
||||
// CountResponse wraps the results of a synchronous batch request to POST /v1/count.
|
||||
type CountResponse struct {
|
||||
Results []PDFResult `json:"results"`
|
||||
Processed int `json:"processed"`
|
||||
Failed int `json:"failed"`
|
||||
}
|
||||
|
||||
// JobStatus tracks async job progress.
|
||||
type JobStatus struct {
|
||||
JobID string `json:"job_id"`
|
||||
Status string `json:"status"`
|
||||
Total int `json:"total"`
|
||||
Processed int `json:"processed"`
|
||||
Failed int `json:"failed"`
|
||||
}
|
||||
|
||||
// JobResultsPage is a paginated slice of PDF results.
|
||||
type JobResultsPage struct {
|
||||
JobID string `json:"job_id"`
|
||||
Results []PDFResult `json:"results"`
|
||||
NextCursor string `json:"next_cursor,omitempty"`
|
||||
HasMore bool `json:"has_more"`
|
||||
}
|
||||
|
||||
// QueueMessage is published per PDF for workers.
|
||||
type QueueMessage struct {
|
||||
JobID string `json:"job_id"`
|
||||
ObjectKey string `json:"object_key"`
|
||||
Filename string `json:"filename"`
|
||||
IncludeRejected bool `json:"include_rejected"`
|
||||
}
|
||||
|
||||
const (
|
||||
StatusOK = "ok"
|
||||
StatusError = "error"
|
||||
|
||||
JobQueued = "queued"
|
||||
JobRunning = "running"
|
||||
JobCompleted = "completed"
|
||||
JobCancelled = "cancelled"
|
||||
)
|
||||
|
||||
@@ -73,6 +73,21 @@ func (p *Processor) ProcessPDF(filename string, pdfData []byte) domain.PDFResult
|
||||
return result
|
||||
}
|
||||
|
||||
// Sin modelo ML: contar todas las filas agrupadas como servicios (heurística).
|
||||
if p.predictor == nil {
|
||||
services := make([]domain.ServiceRow, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
services = append(services, domain.ServiceRow{
|
||||
Text: row.TextoFila,
|
||||
Probability: 1.0,
|
||||
})
|
||||
}
|
||||
result.Services = services
|
||||
result.ServiceCount = len(rows)
|
||||
p.logTiming(filename, renderDur, ocrDur, clusterDur, 0, time.Since(start))
|
||||
return result
|
||||
}
|
||||
|
||||
features := make([][8]float32, len(rows))
|
||||
for i, r := range rows {
|
||||
features[i] = r.Features
|
||||
|
||||
@@ -1,109 +0,0 @@
|
||||
package queue
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/nats-io/nats.go"
|
||||
"github.com/nats-io/nats.go/jetstream"
|
||||
"github.com/pys/pdf-count/internal/domain"
|
||||
)
|
||||
|
||||
type Publisher struct {
|
||||
js jetstream.JetStream
|
||||
stream string
|
||||
subject string
|
||||
}
|
||||
|
||||
func NewPublisher(natsURL, streamName, subjectPrefix string) (*Publisher, error) {
|
||||
nc, err := nats.Connect(natsURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("nats connect: %w", err)
|
||||
}
|
||||
js, err := jetstream.New(nc)
|
||||
if err != nil {
|
||||
nc.Close()
|
||||
return nil, err
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
_, err = js.CreateOrUpdateStream(ctx, jetstream.StreamConfig{
|
||||
Name: streamName,
|
||||
Subjects: []string{subjectPrefix + ".>"},
|
||||
Retention: jetstream.WorkQueuePolicy,
|
||||
MaxMsgs: 200000,
|
||||
})
|
||||
if err != nil {
|
||||
nc.Close()
|
||||
return nil, fmt.Errorf("create stream: %w", err)
|
||||
}
|
||||
|
||||
return &Publisher{js: js, stream: streamName, subject: subjectPrefix}, nil
|
||||
}
|
||||
|
||||
func (p *Publisher) Publish(ctx context.Context, msg domain.QueueMessage) error {
|
||||
data, err := json.Marshal(msg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
subject := fmt.Sprintf("%s.%s", p.subject, msg.JobID)
|
||||
_, err = p.js.Publish(ctx, subject, data)
|
||||
return err
|
||||
}
|
||||
|
||||
type Consumer struct {
|
||||
js jetstream.JetStream
|
||||
subject string
|
||||
consumerName string
|
||||
}
|
||||
|
||||
func NewConsumer(natsURL, streamName, subjectPrefix, consumerName string) (*Consumer, jetstream.Consumer, *nats.Conn, error) {
|
||||
nc, err := nats.Connect(natsURL)
|
||||
if err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
js, err := jetstream.New(nc)
|
||||
if err != nil {
|
||||
nc.Close()
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
_, err = js.CreateOrUpdateStream(ctx, jetstream.StreamConfig{
|
||||
Name: streamName,
|
||||
Subjects: []string{subjectPrefix + ".>"},
|
||||
Retention: jetstream.WorkQueuePolicy,
|
||||
MaxMsgs: 200000,
|
||||
})
|
||||
if err != nil {
|
||||
nc.Close()
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
|
||||
cons, err := js.CreateOrUpdateConsumer(ctx, streamName, jetstream.ConsumerConfig{
|
||||
Name: consumerName,
|
||||
FilterSubject: subjectPrefix + ".>",
|
||||
AckPolicy: jetstream.AckExplicitPolicy,
|
||||
MaxDeliver: 3,
|
||||
AckWait: 2 * time.Minute,
|
||||
MaxAckPending: 256,
|
||||
})
|
||||
if err != nil {
|
||||
nc.Close()
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
|
||||
return &Consumer{js: js, subject: subjectPrefix, consumerName: consumerName}, cons, nc, nil
|
||||
}
|
||||
|
||||
func DecodeMessage(data []byte) (domain.QueueMessage, error) {
|
||||
var msg domain.QueueMessage
|
||||
err := json.Unmarshal(data, &msg)
|
||||
return msg, err
|
||||
}
|
||||
@@ -1,86 +0,0 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
|
||||
"github.com/minio/minio-go/v7"
|
||||
"github.com/minio/minio-go/v7/pkg/credentials"
|
||||
)
|
||||
|
||||
type Client struct {
|
||||
client *minio.Client
|
||||
bucket string
|
||||
}
|
||||
|
||||
func New(endpoint, accessKey, secretKey, bucket string, useSSL bool) (*Client, error) {
|
||||
client, err := minio.New(endpoint, &minio.Options{
|
||||
Creds: credentials.NewStaticV4(accessKey, secretKey, ""),
|
||||
Secure: useSSL,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &Client{client: client, bucket: bucket}, nil
|
||||
}
|
||||
|
||||
func (c *Client) EnsureBucket(ctx context.Context) error {
|
||||
exists, err := c.client.BucketExists(ctx, c.bucket)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if exists {
|
||||
return nil
|
||||
}
|
||||
err = c.client.MakeBucket(ctx, c.bucket, minio.MakeBucketOptions{})
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
resp := minio.ToErrorResponse(err)
|
||||
if resp.Code == "BucketAlreadyOwnedByYou" || resp.Code == "BucketAlreadyExists" {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *Client) Put(ctx context.Context, objectKey string, data []byte, contentType string) error {
|
||||
if contentType == "" {
|
||||
contentType = "application/pdf"
|
||||
}
|
||||
_, err := c.client.PutObject(ctx, c.bucket, objectKey, bytes.NewReader(data), int64(len(data)), minio.PutObjectOptions{
|
||||
ContentType: contentType,
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *Client) Get(ctx context.Context, objectKey string) ([]byte, error) {
|
||||
obj, err := c.client.GetObject(ctx, c.bucket, objectKey, minio.GetObjectOptions{})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer obj.Close()
|
||||
return io.ReadAll(obj)
|
||||
}
|
||||
|
||||
func (c *Client) ListByPrefix(ctx context.Context, prefix string) ([]string, error) {
|
||||
var keys []string
|
||||
for obj := range c.client.ListObjects(ctx, c.bucket, minio.ListObjectsOptions{
|
||||
Prefix: prefix,
|
||||
Recursive: true,
|
||||
}) {
|
||||
if obj.Err != nil {
|
||||
return nil, obj.Err
|
||||
}
|
||||
if strings.HasSuffix(strings.ToLower(obj.Key), ".pdf") {
|
||||
keys = append(keys, obj.Key)
|
||||
}
|
||||
}
|
||||
return keys, nil
|
||||
}
|
||||
|
||||
func JobObjectKey(jobID, filename string) string {
|
||||
return fmt.Sprintf("jobs/%s/%s", jobID, filename)
|
||||
}
|
||||
@@ -1,276 +0,0 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/pys/pdf-count/internal/domain"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
type Store struct {
|
||||
client *redis.Client
|
||||
ttl time.Duration
|
||||
}
|
||||
|
||||
func New(redisURL string, ttl time.Duration) (*Store, error) {
|
||||
opts, err := redis.ParseURL(redisURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
client := redis.NewClient(opts)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
if err := client.Ping(ctx).Err(); err != nil {
|
||||
return nil, fmt.Errorf("redis ping: %w", err)
|
||||
}
|
||||
return &Store{client: client, ttl: ttl}, nil
|
||||
}
|
||||
|
||||
func (s *Store) Close() error {
|
||||
return s.client.Close()
|
||||
}
|
||||
|
||||
func metaKey(jobID string) string { return fmt.Sprintf("job:%s:meta", jobID) }
|
||||
func filesKey(jobID string) string { return fmt.Sprintf("job:%s:files", jobID) }
|
||||
func resultsKey(jobID string) string {
|
||||
return fmt.Sprintf("job:%s:results", jobID)
|
||||
}
|
||||
|
||||
func (s *Store) CreateJob(ctx context.Context, jobID string, total int, includeRejected bool) error {
|
||||
pipe := s.client.Pipeline()
|
||||
pipe.HSet(ctx, metaKey(jobID), map[string]interface{}{
|
||||
"status": domain.JobQueued,
|
||||
"total": total,
|
||||
"processed": 0,
|
||||
"failed": 0,
|
||||
"include_rejected": boolStr(includeRejected),
|
||||
})
|
||||
pipe.Expire(ctx, metaKey(jobID), s.ttl)
|
||||
_, err := pipe.Exec(ctx)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Store) AddJobFiles(ctx context.Context, jobID string, filenames []string) error {
|
||||
if len(filenames) == 0 {
|
||||
return nil
|
||||
}
|
||||
members := make([]interface{}, len(filenames))
|
||||
for i, f := range filenames {
|
||||
members[i] = f
|
||||
}
|
||||
pipe := s.client.Pipeline()
|
||||
pipe.RPush(ctx, filesKey(jobID), members...)
|
||||
pipe.Expire(ctx, filesKey(jobID), s.ttl)
|
||||
_, err := pipe.Exec(ctx)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Store) MarkRunning(ctx context.Context, jobID string) error {
|
||||
return s.client.HSet(ctx, metaKey(jobID), "status", domain.JobRunning).Err()
|
||||
}
|
||||
|
||||
var saveResultScript = redis.NewScript(`
|
||||
local meta = KEYS[1]
|
||||
local results = KEYS[2]
|
||||
redis.call('HSET', results, ARGV[1], ARGV[2])
|
||||
local processed = redis.call('HINCRBY', meta, 'processed', 1)
|
||||
if ARGV[3] == '1' then
|
||||
redis.call('HINCRBY', meta, 'failed', 1)
|
||||
end
|
||||
redis.call('EXPIRE', results, ARGV[4])
|
||||
local total = tonumber(redis.call('HGET', meta, 'total') or '0')
|
||||
local status = redis.call('HGET', meta, 'status')
|
||||
if status ~= 'cancelled' and status ~= 'completed' and total > 0 and processed >= total then
|
||||
redis.call('HSET', meta, 'status', 'completed')
|
||||
end
|
||||
return processed
|
||||
`)
|
||||
|
||||
func (s *Store) SaveResult(ctx context.Context, jobID string, result domain.PDFResult) error {
|
||||
data, err := json.Marshal(result)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
failFlag := "0"
|
||||
if result.Status == domain.StatusError {
|
||||
failFlag = "1"
|
||||
}
|
||||
ttlSec := int64(s.ttl.Seconds())
|
||||
if ttlSec <= 0 {
|
||||
ttlSec = 72 * 3600
|
||||
}
|
||||
_, err = saveResultScript.Run(ctx, s.client,
|
||||
[]string{metaKey(jobID), resultsKey(jobID)},
|
||||
result.Filename, string(data), failFlag, ttlSec,
|
||||
).Result()
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Store) TryCompleteJob(ctx context.Context, jobID string) error {
|
||||
meta, err := s.client.HGetAll(ctx, metaKey(jobID)).Result()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = s.maybeCompleteJob(ctx, jobID, meta)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Store) maybeCompleteJob(ctx context.Context, jobID string, meta map[string]string) (string, error) {
|
||||
status := meta["status"]
|
||||
if status == domain.JobCancelled || status == domain.JobCompleted {
|
||||
return status, nil
|
||||
}
|
||||
total, _ := strconv.Atoi(meta["total"])
|
||||
processed, _ := strconv.Atoi(meta["processed"])
|
||||
if processed >= total && total > 0 {
|
||||
if err := s.client.HSet(ctx, metaKey(jobID), "status", domain.JobCompleted).Err(); err != nil {
|
||||
return status, err
|
||||
}
|
||||
return domain.JobCompleted, nil
|
||||
}
|
||||
return status, nil
|
||||
}
|
||||
|
||||
func (s *Store) GetJob(ctx context.Context, jobID string) (domain.JobStatus, error) {
|
||||
meta, err := s.client.HGetAll(ctx, metaKey(jobID)).Result()
|
||||
if err != nil {
|
||||
return domain.JobStatus{}, err
|
||||
}
|
||||
if len(meta) == 0 {
|
||||
return domain.JobStatus{}, fmt.Errorf("job not found")
|
||||
}
|
||||
status, err := s.maybeCompleteJob(ctx, jobID, meta)
|
||||
if err != nil {
|
||||
return domain.JobStatus{}, err
|
||||
}
|
||||
if status == domain.JobCompleted {
|
||||
meta["status"] = domain.JobCompleted
|
||||
}
|
||||
total, _ := strconv.Atoi(meta["total"])
|
||||
processed, _ := strconv.Atoi(meta["processed"])
|
||||
failed, _ := strconv.Atoi(meta["failed"])
|
||||
return domain.JobStatus{
|
||||
JobID: jobID,
|
||||
Status: status,
|
||||
Total: total,
|
||||
Processed: processed,
|
||||
Failed: failed,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *Store) GetResultsPage(ctx context.Context, jobID, cursor string, limit int) (domain.JobResultsPage, error) {
|
||||
if limit <= 0 || limit > 1000 {
|
||||
limit = 500
|
||||
}
|
||||
allFiles, err := s.client.LRange(ctx, filesKey(jobID), 0, -1).Result()
|
||||
if err != nil {
|
||||
return domain.JobResultsPage{}, err
|
||||
}
|
||||
|
||||
start := 0
|
||||
if cursor != "" {
|
||||
for i, f := range allFiles {
|
||||
if f == cursor {
|
||||
start = i + 1
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
end := start + limit
|
||||
if end > len(allFiles) {
|
||||
end = len(allFiles)
|
||||
}
|
||||
pageFiles := allFiles[start:end]
|
||||
|
||||
results := make([]domain.PDFResult, 0, len(pageFiles))
|
||||
if len(pageFiles) > 0 {
|
||||
raw, err := s.client.HMGet(ctx, resultsKey(jobID), pageFiles...).Result()
|
||||
if err != nil {
|
||||
return domain.JobResultsPage{}, err
|
||||
}
|
||||
for i, v := range raw {
|
||||
if v == nil {
|
||||
results = append(results, domain.PDFResult{
|
||||
Filename: pageFiles[i],
|
||||
Status: "pending",
|
||||
})
|
||||
continue
|
||||
}
|
||||
var r domain.PDFResult
|
||||
if err := json.Unmarshal([]byte(v.(string)), &r); err != nil {
|
||||
continue
|
||||
}
|
||||
results = append(results, r)
|
||||
}
|
||||
}
|
||||
|
||||
page := domain.JobResultsPage{
|
||||
JobID: jobID,
|
||||
Results: results,
|
||||
HasMore: end < len(allFiles),
|
||||
}
|
||||
if page.HasMore && len(pageFiles) > 0 {
|
||||
page.NextCursor = pageFiles[len(pageFiles)-1]
|
||||
}
|
||||
return page, nil
|
||||
}
|
||||
|
||||
func (s *Store) CancelJob(ctx context.Context, jobID string) error {
|
||||
return s.client.HSet(ctx, metaKey(jobID), "status", domain.JobCancelled).Err()
|
||||
}
|
||||
|
||||
func (s *Store) IsCancelled(ctx context.Context, jobID string) (bool, error) {
|
||||
status, err := s.client.HGet(ctx, metaKey(jobID), "status").Result()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return status == domain.JobCancelled, nil
|
||||
}
|
||||
|
||||
func (s *Store) PendingCount(ctx context.Context) (int64, error) {
|
||||
// Approximate: sum of (total - processed) for running jobs
|
||||
var cursor uint64
|
||||
var pending int64
|
||||
for {
|
||||
keys, next, err := s.client.Scan(ctx, cursor, "job:*:meta", 100).Result()
|
||||
if err != nil {
|
||||
return pending, err
|
||||
}
|
||||
for _, key := range keys {
|
||||
meta, err := s.client.HGetAll(ctx, key).Result()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if meta["status"] != domain.JobQueued && meta["status"] != domain.JobRunning {
|
||||
continue
|
||||
}
|
||||
total, _ := strconv.Atoi(meta["total"])
|
||||
processed, _ := strconv.Atoi(meta["processed"])
|
||||
pending += int64(total - processed)
|
||||
}
|
||||
cursor = next
|
||||
if cursor == 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
return pending, nil
|
||||
}
|
||||
|
||||
func (s *Store) IncrementTotal(ctx context.Context, jobID string, delta int) error {
|
||||
pipe := s.client.Pipeline()
|
||||
pipe.HIncrBy(ctx, metaKey(jobID), "total", int64(delta))
|
||||
pipe.Expire(ctx, metaKey(jobID), s.ttl)
|
||||
_, err := pipe.Exec(ctx)
|
||||
return err
|
||||
}
|
||||
|
||||
func boolStr(v bool) string {
|
||||
if v {
|
||||
return "true"
|
||||
}
|
||||
return "false"
|
||||
}
|
||||
@@ -1,212 +0,0 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"os/signal"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/nats-io/nats.go/jetstream"
|
||||
"github.com/pys/pdf-count/internal/config"
|
||||
"github.com/pys/pdf-count/internal/domain"
|
||||
"github.com/pys/pdf-count/internal/model"
|
||||
"github.com/pys/pdf-count/internal/ocr"
|
||||
"github.com/pys/pdf-count/internal/pipeline"
|
||||
"github.com/pys/pdf-count/internal/queue"
|
||||
"github.com/pys/pdf-count/internal/storage"
|
||||
"github.com/pys/pdf-count/internal/store"
|
||||
)
|
||||
|
||||
type Runner struct {
|
||||
cfg config.Config
|
||||
predictor *model.Predictor
|
||||
ocrEngine *ocr.Engine
|
||||
store *store.Store
|
||||
storage *storage.Client
|
||||
consumer jetstream.Consumer
|
||||
ncClose func()
|
||||
}
|
||||
|
||||
func NewRunner(cfg config.Config) (*Runner, error) {
|
||||
predictor, err := model.NewPredictor(cfg.ModelPath, cfg.InferencePath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
poolSize := cfg.OCRPoolSize
|
||||
if poolSize <= 0 {
|
||||
poolSize = cfg.WorkerConcurrency
|
||||
}
|
||||
ocrEngine, err := ocr.NewEngine(cfg.OCRLang, poolSize)
|
||||
if err != nil {
|
||||
predictor.Close()
|
||||
return nil, err
|
||||
}
|
||||
st, err := store.New(cfg.RedisURL, cfg.JobResultTTL)
|
||||
if err != nil {
|
||||
ocrEngine.Close()
|
||||
predictor.Close()
|
||||
return nil, err
|
||||
}
|
||||
stg, err := storage.New(cfg.MinIOEndpoint, cfg.MinIOAccessKey, cfg.MinIOSecretKey, cfg.MinIOBucket, cfg.MinIOUseSSL)
|
||||
if err != nil {
|
||||
ocrEngine.Close()
|
||||
st.Close()
|
||||
predictor.Close()
|
||||
return nil, err
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
if err := stg.EnsureBucket(ctx); err != nil {
|
||||
ocrEngine.Close()
|
||||
st.Close()
|
||||
predictor.Close()
|
||||
return nil, err
|
||||
}
|
||||
|
||||
_, cons, nc, err := queue.NewConsumer(cfg.NATSURL, cfg.StreamName, cfg.SubjectPrefix, "pdf-workers")
|
||||
if err != nil {
|
||||
ocrEngine.Close()
|
||||
st.Close()
|
||||
predictor.Close()
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &Runner{
|
||||
cfg: cfg,
|
||||
predictor: predictor,
|
||||
ocrEngine: ocrEngine,
|
||||
store: st,
|
||||
storage: stg,
|
||||
consumer: cons,
|
||||
ncClose: func() { nc.Close() },
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *Runner) Close() {
|
||||
r.ncClose()
|
||||
_ = r.ocrEngine.Close()
|
||||
_ = r.store.Close()
|
||||
_ = r.predictor.Close()
|
||||
}
|
||||
|
||||
func (r *Runner) processSlots() int {
|
||||
c := r.cfg.WorkerConcurrency
|
||||
pool := r.cfg.OCRPoolSize
|
||||
if pool <= 0 {
|
||||
pool = r.cfg.WorkerConcurrency
|
||||
}
|
||||
if pool <= 0 {
|
||||
pool = 4
|
||||
}
|
||||
if c > pool {
|
||||
c = pool
|
||||
}
|
||||
if c <= 0 {
|
||||
c = 4
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
func (r *Runner) Run(ctx context.Context) error {
|
||||
slots := r.processSlots()
|
||||
sem := make(chan struct{}, slots)
|
||||
var wg sync.WaitGroup
|
||||
|
||||
go func() {
|
||||
<-ctx.Done()
|
||||
wg.Wait()
|
||||
}()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil
|
||||
default:
|
||||
}
|
||||
|
||||
msgs, err := r.consumer.Fetch(slots, jetstream.FetchMaxWait(r.cfg.NATSFetchMaxWait))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
for msg := range msgs.Messages() {
|
||||
sem <- struct{}{}
|
||||
wg.Add(1)
|
||||
m := msg
|
||||
go func() {
|
||||
defer func() {
|
||||
<-sem
|
||||
wg.Done()
|
||||
}()
|
||||
r.handleMessage(ctx, m)
|
||||
}()
|
||||
}
|
||||
_ = msgs.Error()
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Runner) handleMessage(ctx context.Context, msg jetstream.Msg) {
|
||||
processCtx, cancel := context.WithTimeout(ctx, r.cfg.PDFProcessTimeout)
|
||||
defer cancel()
|
||||
|
||||
qmsg, err := queue.DecodeMessage(msg.Data())
|
||||
if err != nil {
|
||||
_ = msg.Term()
|
||||
return
|
||||
}
|
||||
|
||||
cancelled, err := r.store.IsCancelled(processCtx, qmsg.JobID)
|
||||
if err != nil || cancelled {
|
||||
_ = msg.Ack()
|
||||
return
|
||||
}
|
||||
|
||||
data, err := r.storage.Get(processCtx, qmsg.ObjectKey)
|
||||
if err != nil {
|
||||
_ = r.store.SaveResult(processCtx, qmsg.JobID, domain.PDFResult{
|
||||
Filename: qmsg.Filename,
|
||||
Status: domain.StatusError,
|
||||
Error: err.Error(),
|
||||
})
|
||||
_ = r.store.TryCompleteJob(processCtx, qmsg.JobID)
|
||||
_ = msg.Ack()
|
||||
return
|
||||
}
|
||||
|
||||
proc := pipeline.NewProcessor(r.predictor, r.ocrEngine, pipeline.Options{
|
||||
DPI: r.cfg.OCRDPI,
|
||||
Lang: r.cfg.OCRLang,
|
||||
DBSCANEps: r.cfg.DBSCANEps,
|
||||
DBSCANMinSamples: r.cfg.DBSCANMinSamples,
|
||||
IncludeRejected: qmsg.IncludeRejected,
|
||||
LogPipelineTiming: r.cfg.LogPipelineTiming,
|
||||
})
|
||||
result := proc.ProcessPDF(qmsg.Filename, data)
|
||||
data = nil
|
||||
|
||||
if err := r.store.SaveResult(processCtx, qmsg.JobID, result); err != nil {
|
||||
_ = msg.NakWithDelay(5 * time.Second)
|
||||
return
|
||||
}
|
||||
_ = r.store.TryCompleteJob(processCtx, qmsg.JobID)
|
||||
_ = msg.Ack()
|
||||
}
|
||||
|
||||
func RunMain() {
|
||||
cfg := config.Load()
|
||||
runner, err := NewRunner(cfg)
|
||||
if err != nil {
|
||||
log.Fatalf("worker init: %v", err)
|
||||
}
|
||||
defer runner.Close()
|
||||
|
||||
ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
||||
defer cancel()
|
||||
|
||||
log.Printf("worker started slots=%d", runner.processSlots())
|
||||
if err := runner.Run(ctx); err != nil {
|
||||
log.Fatalf("worker: %v", err)
|
||||
}
|
||||
log.Println("worker stopped")
|
||||
}
|
||||
Reference in New Issue
Block a user