agent SDK fix

This commit is contained in:
m.nazemi
2026-04-12 12:29:35 +03:30
parent fd89371bdf
commit bf3f21fb51
8 changed files with 1186 additions and 20 deletions
+84 -16
View File
@@ -13,6 +13,7 @@ import (
"tm/pkg/mongo"
"tm/pkg/notification"
"tm/pkg/ollama"
"tm/pkg/queue"
"tm/pkg/schedule"
"tm/pkg/scraper"
)
@@ -81,6 +82,8 @@ func InitWorker(config Config, mongoManager *mongo.ConnectionManager, appLogger
// Debug: Log worker config
appLogger.Info("Worker configuration", map[string]interface{}{
"worker_interval": config.Worker.Interval,
"queue_enabled": config.Queue.Enabled,
"queue_mode": config.Queue.Mode,
})
// Initialize repositories
noticeRepo := notice.NewRepository(mongoManager, appLogger)
@@ -91,11 +94,87 @@ func InitWorker(config Config, mongoManager *mongo.ConnectionManager, appLogger
noticeWorker := workers.NewNoticeWorker(mongoManager, appLogger, &notify, noticeRepo, tenderRepo, glmService, scraperService)
noticeWorker.Run()
// Initialize Document Scraper Worker
// appLogger.Info("Initializing document scraper worker", map[string]interface{}{})
// scraperWorker := workers.NewDocumentScraperWorker(mongoManager, appLogger, tenderRepo, scraperService)
// appLogger.Info("Starting document scraper worker", map[string]interface{}{})
// scraperWorker.Run()
// Initialize Queue-based scraping if enabled
var queueService queue.Queue
if config.Queue.Enabled {
var err error
queueConfig := queue.QueueConfig{
RedisURL: config.Queue.RedisURL,
MaxRetries: config.Queue.MaxRetries,
RetryBackoff: config.Queue.RetryBackoff,
JobTimeout: config.Queue.JobTimeout,
PoolSize: config.Queue.PoolSize,
IsDLQEnabled: config.Queue.IsDLQEnabled,
MaxQueueLength: 0, // Unlimited
}
queueService, err = queue.NewRedisQueue(queueConfig, appLogger)
if err != nil {
appLogger.Error("Failed to initialize queue service", map[string]interface{}{
"error": err.Error(),
})
appLogger.Warn("Queue-based scraping is disabled, falling back to scheduled scraping", map[string]interface{}{})
config.Queue.Enabled = false
} else {
appLogger.Info("Queue service initialized successfully", map[string]interface{}{
"mode": config.Queue.Mode,
})
}
}
// Start Queue Consumer Worker (24/7 continuous processing)
if config.Queue.Enabled && (config.Queue.Mode == "consumer" || config.Queue.Mode == "both") {
appLogger.Info("Starting queue consumer worker", map[string]interface{}{
"concurrency": config.Queue.Concurrency,
})
consumerWorker := workers.NewDocumentScraperConsumerWorker(
mongoManager,
appLogger,
tenderRepo,
queueService,
scraperService,
config.Queue.Concurrency,
)
go consumerWorker.Run()
} else if !config.Queue.Enabled {
appLogger.Info("Queue consumer worker disabled, using scheduled scraping instead", map[string]interface{}{})
}
// Start Queue Producer Worker (enqueues unscraped tenders periodically)
if config.Queue.Enabled && (config.Queue.Mode == "producer" || config.Queue.Mode == "both") {
appLogger.Info("Scheduling queue producer worker", map[string]interface{}{
"batch_size": config.Queue.ProducerBatch,
"interval": config.Queue.ProducerInterval,
})
schedule.NewCronScheduler(appLogger, mongoManager, noticeRepo).AddJob(schedule.Job{
Name: "Queue Producer Worker Job",
Func: func() {
producer := workers.NewQueueProducerWorker(
mongoManager,
appLogger,
tenderRepo,
queueService,
config.Queue.ProducerBatch,
)
producer.Run()
},
Expr: config.Queue.ProducerInterval,
})
}
// Fallback: Keep the old scheduled scraper for backward compatibility
if !config.Queue.Enabled {
appLogger.Info("Scheduling fallback document scraper worker", map[string]interface{}{})
schedule.NewCronScheduler(appLogger, mongoManager, noticeRepo).AddJob(schedule.Job{
Name: "Document Scraper Worker Job (Fallback)",
Func: func() {
worker := workers.NewDocumentScraperWorker(mongoManager, appLogger, tenderRepo, scraperService)
worker.Run()
},
Expr: "0 10 * * * *", // Run at 10 AM every day
})
}
// Initialize Document Summarization Worker (only if MinIO is available)
if minioService != nil {
@@ -133,17 +212,6 @@ func InitWorker(config Config, mongoManager *mongo.ConnectionManager, appLogger
},
Expr: workerInterval,
})
// Schedule Document Scraper Worker job (runs after notice worker)
schedule.NewCronScheduler(appLogger, mongoManager, noticeRepo).AddJob(schedule.Job{
Name: "Document Scraper Worker Job",
Func: func() {
worker := workers.NewDocumentScraperWorker(mongoManager, appLogger, tenderRepo, scraperService)
worker.Run()
},
Expr: "0 10 * * * *", // Run at 10 AM every day (1 hour after notice worker)
})
}
// Init Notification Service
+16
View File
@@ -15,9 +15,25 @@ type Config struct {
GLM GLMConfig
Scraper config.ScraperConfig
MinIO config.MinIOConfig
Queue QueueConfig `env:"QUEUE_" envPrefix:"QUEUE_"`
AlertMail string `env:"ALERT_MAIL" envDefault:"alerts@opplens.com"`
}
// QueueConfig defines the queue configuration
type QueueConfig struct {
Enabled bool `env:"QUEUE_ENABLED" envDefault:"true"`
RedisURL string `env:"QUEUE_REDIS_URL" envDefault:"redis://localhost:6379"`
Mode string `env:"QUEUE_MODE" envDefault:"consumer"` // "producer", "consumer", or "both"
MaxRetries int `env:"QUEUE_MAX_RETRIES" envDefault:"3"`
RetryBackoff time.Duration `env:"QUEUE_RETRY_BACKOFF" envDefault:"10s"`
JobTimeout time.Duration `env:"QUEUE_JOB_TIMEOUT" envDefault:"5m"`
PoolSize int `env:"QUEUE_POOL_SIZE" envDefault:"10"`
IsDLQEnabled bool `env:"QUEUE_DLQ_ENABLED" envDefault:"true"`
Concurrency int `env:"QUEUE_CONCURRENCY" envDefault:"5"`
ProducerBatch int `env:"QUEUE_PRODUCER_BATCH" envDefault:"100"`
ProducerInterval string `env:"QUEUE_PRODUCER_INTERVAL" envDefault:"0 */6 * * * *"` // Every 6 hours
}
type WorkerConfig struct {
Interval string `env:"WORKER_INTERVAL" envDefault:"* 10 * * * *"`
}
+3 -3
View File
@@ -573,13 +573,13 @@ func (w *NoticeWorker) getNextSequentialNumber(ctx context.Context, sourceCode,
// Create aggregation pipeline to find the highest sequential number
pipeline := mongodriver.Pipeline{
bson.D{{"$match", bson.M{
bson.D{{Key: "$match", Value: bson.M{
"tender_id": bson.M{"$regex": pattern},
}}},
bson.D{{"$project", bson.M{
bson.D{{Key: "$project", Value: bson.M{
"sequential_num": bson.M{"$substr": bson.A{"$tender_id", 5, 3}}, // Extract last 3 chars (NNN)
}}},
bson.D{{"$group", bson.M{
bson.D{{Key: "$group", Value: bson.M{
"_id": nil,
"max_num": bson.M{"$max": bson.M{"$toInt": "$sequential_num"}},
}}},
+165
View File
@@ -0,0 +1,165 @@
package workers
import (
"context"
"time"
"tm/internal/tender"
"tm/pkg/logger"
"tm/pkg/mongo"
"tm/pkg/queue"
)
// QueueProducerWorker handles enqueuing unscraped tenders to the job queue
type QueueProducerWorker struct {
Mongo *mongo.ConnectionManager
Logger logger.Logger
TenderRepo tender.TenderRepository
Queue queue.Queue
BatchSize int
}
// NewQueueProducerWorker creates a new queue producer worker
func NewQueueProducerWorker(
mongo *mongo.ConnectionManager,
logger logger.Logger,
tenderRepo tender.TenderRepository,
queue queue.Queue,
batchSize int,
) *QueueProducerWorker {
if batchSize <= 0 {
batchSize = 100 // Default batch size
}
logger.Info("Creating new queue producer worker", map[string]interface{}{
"batch_size": batchSize,
})
return &QueueProducerWorker{
Mongo: mongo,
Logger: logger,
TenderRepo: tenderRepo,
Queue: queue,
BatchSize: batchSize,
}
}
// Run continuously processes and enqueues unscraped tenders
func (w *QueueProducerWorker) Run() {
w.Logger.Info("Queue producer worker started", map[string]interface{}{})
defer func() {
w.Logger.Info("Queue producer worker stopped", map[string]interface{}{})
}()
batchIndex := 0
hasMore := true
for hasMore {
w.Logger.Info("Queue producer fetching batch of tenders", map[string]interface{}{
"batch_index": batchIndex,
"batch_size": w.BatchSize,
})
// Get batch of unscraped tenders
tenders, totalCount, err := w.TenderRepo.GetUnScrapedTenders(
context.Background(),
w.BatchSize,
batchIndex*w.BatchSize,
)
if err != nil {
w.Logger.Error("Failed to fetch unscraped tenders", map[string]interface{}{
"batch_index": batchIndex,
"error": err.Error(),
})
// Wait before retrying
time.Sleep(30 * time.Second)
continue
}
w.Logger.Info("Fetched tender batch", map[string]interface{}{
"batch_size": len(tenders),
"total_count": totalCount,
"batch_index": batchIndex,
})
if len(tenders) == 0 {
// No more tenders in this batch
hasMore = false
w.Logger.Info("No more unscraped tenders found", map[string]interface{}{
"total_processed": totalCount,
})
break
}
// Enqueue each tender
successCount := 0
failCount := 0
for _, t := range tenders {
if err := w.enqueueTender(&t); err != nil {
w.Logger.Error("Failed to enqueue tender", map[string]interface{}{
"tender_id": t.ID.Hex(),
"notice_publication_id": t.NoticePublicationID,
"error": err.Error(),
})
failCount++
} else {
successCount++
}
}
w.Logger.Info("Batch enqueuing completed", map[string]interface{}{
"success_count": successCount,
"fail_count": failCount,
"batch_index": batchIndex,
})
// Move to next batch
batchIndex++
// Check if there are more tenders
if int64(batchIndex*w.BatchSize) >= totalCount {
hasMore = false
}
// Avoid tight loop - brief pause between batches
time.Sleep(1 * time.Second)
}
}
// enqueueTender adds a single tender to the queue
func (w *QueueProducerWorker) enqueueTender(t *tender.Tender) error {
// Skip if missing required fields
if t.NoticePublicationID == "" || t.DocumentURI == "" {
w.Logger.Warn("Skipping tender - missing required fields", map[string]interface{}{
"tender_id": t.ID.Hex(),
"notice_publication_id": t.NoticePublicationID,
"document_uri": t.DocumentURI,
})
return nil
}
// Create job
job := queue.NewDocumentScraperJob(
t.NoticePublicationID,
t.DocumentURI,
t.ID.Hex(),
queue.DefaultMaxRetries,
)
// Enqueue job
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := w.Queue.Enqueue(ctx, job); err != nil {
return err
}
return nil
}
// RunOnce processes all unscraped tenders in a single run (for scheduled use)
func (w *QueueProducerWorker) RunOnce() {
w.Run()
}
+422
View File
@@ -0,0 +1,422 @@
package workers
import (
"context"
"fmt"
"time"
tenderPkg "tm/internal/tender"
"tm/pkg/logger"
"tm/pkg/mongo"
"tm/pkg/queue"
"tm/pkg/scraper"
)
// DocumentScraperConsumerWorker processes scraping jobs from the queue continuously
type DocumentScraperConsumerWorker struct {
Mongo *mongo.ConnectionManager
Logger logger.Logger
TenderRepo tenderPkg.TenderRepository
Queue queue.Queue
Scraper scraper.SDK
Concurrency int
StopChan chan struct{}
}
// NewDocumentScraperConsumerWorker creates a new consumer worker
func NewDocumentScraperConsumerWorker(
mongo *mongo.ConnectionManager,
logger logger.Logger,
tenderRepo tenderPkg.TenderRepository,
queue queue.Queue,
scraperService scraper.SDK,
concurrency int,
) *DocumentScraperConsumerWorker {
if concurrency <= 0 {
concurrency = 5 // Default concurrency
}
logger.Info("Creating new document scraper consumer worker", map[string]interface{}{
"concurrency": concurrency,
})
return &DocumentScraperConsumerWorker{
Mongo: mongo,
Logger: logger,
TenderRepo: tenderRepo,
Queue: queue,
Scraper: scraperService,
Concurrency: concurrency,
StopChan: make(chan struct{}),
}
}
// Run starts consuming jobs from the queue with specified concurrency
func (w *DocumentScraperConsumerWorker) Run() {
w.Logger.Info("Document scraper consumer worker started", map[string]interface{}{
"concurrency": w.Concurrency,
})
// Create worker pool channels
jobsChan := make(chan *queue.DocumentScraperJob, w.Concurrency)
resultsChan := make(chan *ProcessingResult, w.Concurrency)
// Start worker goroutines
for i := 0; i < w.Concurrency; i++ {
go w.processJobs(i, jobsChan, resultsChan)
}
// Start job fetcher
go w.fetchJobs(jobsChan)
// Start result handler
go w.handleResults(resultsChan)
// Wait for stop signal
<-w.StopChan
w.Logger.Info("Document scraper consumer worker stopping", map[string]interface{}{})
close(jobsChan)
close(resultsChan)
}
// Stop signals the worker to stop
func (w *DocumentScraperConsumerWorker) Stop() {
w.Logger.Info("Stop signal sent to consumer worker", map[string]interface{}{})
close(w.StopChan)
}
// fetchJobs continuously fetches jobs from the queue and sends them to workers
func (w *DocumentScraperConsumerWorker) fetchJobs(jobsChan chan<- *queue.DocumentScraperJob) {
defer func() {
w.Logger.Info("Job fetcher stopped", map[string]interface{}{})
}()
for {
select {
case <-w.StopChan:
return
default:
// Try to dequeue a job
ctx, cancel := context.WithTimeout(context.Background(), 35*time.Second)
job, err := w.Queue.Dequeue(ctx)
cancel()
if err != nil {
w.Logger.Error("Failed to dequeue job", map[string]interface{}{
"error": err.Error(),
})
time.Sleep(5 * time.Second) // Back off on error
continue
}
if job == nil {
// Queue is empty, brief wait before retry
time.Sleep(2 * time.Second)
continue
}
// Send to processing channel
select {
case jobsChan <- job:
// Job sent successfully
case <-w.StopChan:
return
}
}
}
}
// processJobs processes jobs from the jobs channel
func (w *DocumentScraperConsumerWorker) processJobs(
workerID int,
jobsChan <-chan *queue.DocumentScraperJob,
resultsChan chan<- *ProcessingResult,
) {
w.Logger.Info("Worker started", map[string]interface{}{
"worker_id": workerID,
})
defer func() {
w.Logger.Info("Worker stopped", map[string]interface{}{
"worker_id": workerID,
})
}()
for {
select {
case job, ok := <-jobsChan:
if !ok {
return // Channel closed
}
if job == nil {
continue
}
// Process the job
result := w.processJob(job, workerID)
select {
case resultsChan <- result:
// Result sent successfully
case <-w.StopChan:
return
}
case <-w.StopChan:
return
}
}
}
// processJob executes the actual scraping task
func (w *DocumentScraperConsumerWorker) processJob(
job *queue.DocumentScraperJob,
workerID int,
) *ProcessingResult {
startTime := time.Now()
w.Logger.Info("Processing job", map[string]interface{}{
"worker_id": workerID,
"job_id": job.ID,
"notice_id": job.NoticeID,
"attempt": job.Attempts,
"max_retries": job.MaxRetries,
})
result := &ProcessingResult{
JobID: job.ID,
WorkerID: workerID,
StartedAt: startTime,
Job: job,
}
// Skip if missing required fields
if job.NoticeID == "" || job.NoticeURL == "" {
result.Success = false
result.Error = "missing required fields: notice_id or notice_url"
result.CompletedAt = time.Now()
return result
}
// Call scraper SDK with timeout
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
defer cancel()
scrapeReq := &scraper.ScrapeRequest{
NoticeID: job.NoticeID,
NoticeURL: job.NoticeURL,
}
w.Logger.Info("Calling scraper service", map[string]interface{}{
"job_id": job.ID,
"notice_id": job.NoticeID,
"notice_url": job.NoticeURL,
"worker_id": workerID,
})
response, err := w.Scraper.ScrapeDocuments(ctx, scrapeReq)
if err != nil {
result.Success = false
result.Error = err.Error()
result.CompletedAt = time.Now()
result.ShouldRetry = job.Attempts < job.MaxRetries
w.Logger.Error("Scraping failed", map[string]interface{}{
"job_id": job.ID,
"worker_id": workerID,
"attempt": job.Attempts,
"error": err.Error(),
"should_retry": result.ShouldRetry,
})
return result
}
if !response.Success {
result.Success = false
result.Error = "scraper returned failure"
if len(response.Errors) > 0 {
result.Error += ": " + response.Errors[0].Error
}
result.CompletedAt = time.Now()
result.ShouldRetry = job.Attempts < job.MaxRetries
w.Logger.Warn("Scraping reported failure", map[string]interface{}{
"job_id": job.ID,
"worker_id": workerID,
"error_count": len(response.Errors),
"should_retry": result.ShouldRetry,
})
return result
}
// Update tender with scraped documents
scrapedAt := time.Now().Unix()
tender, err := w.TenderRepo.GetByNoticePublicationID(context.Background(), job.NoticeID)
if err != nil {
result.Success = false
result.Error = "failed to fetch tender: " + err.Error()
result.CompletedAt = time.Now()
result.ShouldRetry = job.Attempts < job.MaxRetries
w.Logger.Error("Failed to fetch tender for update", map[string]interface{}{
"job_id": job.ID,
"notice_id": job.NoticeID,
"error": err.Error(),
"worker_id": workerID,
})
return result
}
if tender == nil {
result.Success = false
result.Error = "tender not found"
result.CompletedAt = time.Now()
w.Logger.Warn("Tender not found for scraping result", map[string]interface{}{
"job_id": job.ID,
"notice_id": job.NoticeID,
"worker_id": workerID,
})
return result
}
// Update tender with scraped documents
tender.ScrapedDocuments = make([]tenderPkg.ScrapedDocument, len(response.UploadedFiles))
for i, uploadedFile := range response.UploadedFiles {
tender.ScrapedDocuments[i] = tenderPkg.ScrapedDocument{
Filename: uploadedFile.Filename,
ObjectName: uploadedFile.ObjectName,
BucketName: "opplens", // Default bucket
Size: uploadedFile.Size,
ScrapedAt: scrapedAt,
}
}
// Update metadata
tender.ProcessingMetadata.DocumentsScraped = true
tender.ProcessingMetadata.DocumentsScrapedAt = scrapedAt
// Save to database
if err := w.TenderRepo.Update(context.Background(), tender); err != nil {
result.Success = false
result.Error = "failed to update tender: " + err.Error()
result.CompletedAt = time.Now()
result.ShouldRetry = job.Attempts < job.MaxRetries
w.Logger.Error("Failed to update tender with scraped documents", map[string]interface{}{
"job_id": job.ID,
"tender_id": tender.ID.Hex(),
"error": err.Error(),
"worker_id": workerID,
"should_retry": result.ShouldRetry,
})
return result
}
result.Success = true
result.UploadedCount = response.UploadedCount
result.CompletedAt = time.Now()
w.Logger.Info("Scraping job completed successfully", map[string]interface{}{
"job_id": job.ID,
"tender_id": tender.ID.Hex(),
"uploaded_count": response.UploadedCount,
"worker_id": workerID,
"duration_seconds": time.Since(startTime).Seconds(),
})
return result
}
// handleResults processes the results from worker jobs
func (w *DocumentScraperConsumerWorker) handleResults(resultsChan <-chan *ProcessingResult) {
defer func() {
w.Logger.Info("Result handler stopped", map[string]interface{}{})
}()
for {
select {
case result, ok := <-resultsChan:
if !ok {
return // Channel closed
}
if result == nil {
continue
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
if result.Success {
// Acknowledge successful job
if err := w.Queue.AckJob(ctx, result.JobID); err != nil {
w.Logger.Error("Failed to acknowledge job", map[string]interface{}{
"job_id": result.JobID,
"error": err.Error(),
})
}
} else {
// Handle failure - potentially retry
if result.ShouldRetry {
w.Logger.Info("Job will be retried", map[string]interface{}{
"job_id": result.JobID,
"attempt": result.Job.Attempts,
"error": result.Error,
"retry_in": "30 seconds",
})
// Re-enqueue the job for retry
job := result.Job
job.NextRetryTime = time.Now().Add(30 * time.Second).Unix()
job.Status = string(queue.StatusRetrying)
if err := w.Queue.Enqueue(ctx, job); err != nil {
w.Logger.Error("Failed to re-enqueue job for retry", map[string]interface{}{
"job_id": job.ID,
"error": err.Error(),
})
// Still nack it if re-enqueue fails
w.Queue.NackJob(ctx, result.JobID, fmt.Errorf(result.Error))
}
} else {
// Max retries exceeded, send to DLQ
w.Logger.Error("Job exceeded max retries, sending to DLQ", map[string]interface{}{
"job_id": result.JobID,
"max_retry": result.Job.MaxRetries,
"error": result.Error,
})
if err := w.Queue.NackJob(ctx, result.JobID, fmt.Errorf(result.Error)); err != nil {
w.Logger.Error("Failed to nack job", map[string]interface{}{
"job_id": result.JobID,
"error": err.Error(),
})
}
}
}
cancel()
case <-w.StopChan:
return
}
}
}
// ProcessingResult represents the result of processing a job
type ProcessingResult struct {
JobID string
WorkerID int
Success bool
Error string
ShouldRetry bool
UploadedCount int
StartedAt time.Time
CompletedAt time.Time
Job *queue.DocumentScraperJob
}
+53
View File
@@ -0,0 +1,53 @@
package queue
import (
"time"
)
/*
Queue Configuration and Constants
This package provides a Redis-based job queue system for managing
document scraping tasks in a distributed manner.
*/
const (
// Queue names
DocumentScraperQueue = "scraper:document:queue"
DocumentScraperDLQ = "scraper:document:dlq" // Dead Letter Queue
// Job status keys
JobStatusPrefix = "job:"
QueueMetricsKey = "queue:metrics"
// Default timeouts
DefaultJobTimeout = 5 * time.Minute
DefaultMaxRetries = 3
DefaultRetryBackoff = 10 * time.Second
MaxRetryBackoff = 5 * time.Minute // Cap retry backoff at 5 minutes
InProgressExpiration = 10 * time.Minute // Job must complete within this time
)
// QueueConfig represents configuration for the queue service
type QueueConfig struct {
RedisURL string
MaxRetries int
RetryBackoff time.Duration
JobTimeout time.Duration
PoolSize int
IsDLQEnabled bool
MaxQueueLength int // 0 means unlimited
}
// DefaultConfig returns default queue configuration
func DefaultConfig() QueueConfig {
return QueueConfig{
RedisURL: "redis://localhost:6379",
MaxRetries: DefaultMaxRetries,
RetryBackoff: DefaultRetryBackoff,
JobTimeout: DefaultJobTimeout,
PoolSize: 10,
IsDLQEnabled: true,
MaxQueueLength: 0, // Unlimited by default
}
}
+85
View File
@@ -0,0 +1,85 @@
package queue
import (
"encoding/json"
"time"
)
// JobStatus represents the status of a queue job
type JobStatus string
const (
StatusPending JobStatus = "pending"
StatusInProgress JobStatus = "in_progress"
StatusCompleted JobStatus = "completed"
StatusFailed JobStatus = "failed"
StatusRetrying JobStatus = "retrying"
)
// DocumentScraperJob represents a document scraping task in the queue
type DocumentScraperJob struct {
ID string `json:"id"` // Unique job ID (typically tender ID)
NoticeID string `json:"notice_id"` // Notice publication ID
NoticeURL string `json:"notice_url"` // Document URL to scrape
TenderID string `json:"tender_id,omitempty"` // Associated tender ID from MongoDB
CreatedAt int64 `json:"created_at"`
UpdatedAt int64 `json:"updated_at"`
Attempts int `json:"attempts"`
MaxRetries int `json:"max_retries"`
Status string `json:"status"`
Error string `json:"error,omitempty"`
NextRetryTime int64 `json:"next_retry_time,omitempty"`
CompletedAt int64 `json:"completed_at,omitempty"`
ProcessingStartedAt int64 `json:"processing_started_at,omitempty"`
}
// NewDocumentScraperJob creates a new document scraper job
func NewDocumentScraperJob(noticeID, noticeURL, tenderID string, maxRetries int) *DocumentScraperJob {
now := time.Now().Unix()
return &DocumentScraperJob{
ID: noticeID,
NoticeID: noticeID,
NoticeURL: noticeURL,
TenderID: tenderID,
CreatedAt: now,
UpdatedAt: now,
Attempts: 0,
MaxRetries: maxRetries,
Status: string(StatusPending),
}
}
// ToJSON marshals the job to JSON
func (j *DocumentScraperJob) ToJSON() (string, error) {
data, err := json.Marshal(j)
return string(data), err
}
// FromJSON unmarshals JSON to DocumentScraperJob
func (j *DocumentScraperJob) FromJSON(data string) error {
return json.Unmarshal([]byte(data), j)
}
// QueueMetrics represents queue performance metrics
type QueueMetrics struct {
TotalJobs int64 `json:"total_jobs"`
PendingJobs int64 `json:"pending_jobs"`
InProgressJobs int64 `json:"in_progress_jobs"`
CompletedJobs int64 `json:"completed_jobs"`
FailedJobs int64 `json:"failed_jobs"`
AverageWaitTime time.Duration `json:"average_wait_time"`
AverageProcessTime time.Duration `json:"average_process_time"`
UpdatedAt int64 `json:"updated_at"`
}
// JobResult represents the result of a completed job
type JobResult struct {
JobID string `json:"job_id"`
Status string `json:"status"`
Success bool `json:"success"`
Error string `json:"error,omitempty"`
CompletedAt int64 `json:"completed_at"`
Duration int64 `json:"duration_ms"`
UploadedCount int `json:"uploaded_count"`
Details map[string]interface{} `json:"details,omitempty"`
}
+357
View File
@@ -0,0 +1,357 @@
package queue
import (
"context"
"fmt"
"time"
"tm/pkg/logger"
"github.com/redis/go-redis/v9"
)
// Queue interface defines the queue operations
type Queue interface {
// Job operations
Enqueue(ctx context.Context, job *DocumentScraperJob) error
Dequeue(ctx context.Context) (*DocumentScraperJob, error)
AckJob(ctx context.Context, jobID string) error
NackJob(ctx context.Context, jobID string, err error) error
GetJobStatus(ctx context.Context, jobID string) (JobStatus, error)
// Queue inspection
GetMetrics(ctx context.Context) (*QueueMetrics, error)
GetQueueSize(ctx context.Context) (int64, error)
GetDLQSize(ctx context.Context) (int64, error)
GetInProgressJobs(ctx context.Context) ([]DocumentScraperJob, error)
// Cleanup operations
PurgeQueue(ctx context.Context) error
PurgeDLQ(ctx context.Context) error
Close() error
}
// RedisQueue implements Queue interface using Redis
type RedisQueue struct {
client *redis.Client
config QueueConfig
logger logger.Logger
}
// NewRedisQueue creates a new Redis-based queue
func NewRedisQueue(config QueueConfig, logger logger.Logger) (Queue, error) {
opt, err := redis.ParseURL(config.RedisURL)
if err != nil {
logger.Error("Failed to parse Redis URL", map[string]interface{}{
"url": config.RedisURL,
"error": err.Error(),
})
return nil, fmt.Errorf("failed to parse redis url: %w", err)
}
opt.PoolSize = config.PoolSize
client := redis.NewClient(opt)
// Test connection
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := client.Ping(ctx).Err(); err != nil {
logger.Error("Failed to connect to Redis", map[string]interface{}{
"error": err.Error(),
})
return nil, fmt.Errorf("failed to connect to redis: %w", err)
}
logger.Info("Redis queue initialized successfully", map[string]interface{}{
"pool_size": config.PoolSize,
})
return &RedisQueue{
client: client,
config: config,
logger: logger,
}, nil
}
// Enqueue adds a job to the queue
func (q *RedisQueue) Enqueue(ctx context.Context, job *DocumentScraperJob) error {
if q.config.MaxQueueLength > 0 {
size, err := q.GetQueueSize(ctx)
if err != nil {
q.logger.Warn("Failed to check queue size", map[string]interface{}{
"error": err.Error(),
})
} else if size >= int64(q.config.MaxQueueLength) {
return fmt.Errorf("queue is full: %d/%d", size, q.config.MaxQueueLength)
}
}
job.CreatedAt = time.Now().Unix()
job.UpdatedAt = time.Now().Unix()
job.Status = string(StatusPending)
jsonData, err := job.ToJSON()
if err != nil {
q.logger.Error("Failed to marshal job", map[string]interface{}{
"job_id": job.ID,
"error": err.Error(),
})
return err
}
// Push to queue
if err := q.client.RPush(ctx, DocumentScraperQueue, jsonData).Err(); err != nil {
q.logger.Error("Failed to enqueue job", map[string]interface{}{
"job_id": job.ID,
"error": err.Error(),
})
return err
}
// Store job metadata with TTL
jobKey := fmt.Sprintf("%s%s", JobStatusPrefix, job.ID)
if err := q.client.Set(ctx, jobKey, string(StatusPending), 24*time.Hour).Err(); err != nil {
q.logger.Warn("Failed to set job status", map[string]interface{}{
"job_id": job.ID,
"error": err.Error(),
})
}
q.logger.Info("Job enqueued successfully", map[string]interface{}{
"job_id": job.ID,
"notice_id": job.NoticeID,
})
return nil
}
// Dequeue retrieves and locks the next job from the queue
func (q *RedisQueue) Dequeue(ctx context.Context) (*DocumentScraperJob, error) {
// Use BLPOP to get next job with timeout
result, err := q.client.BLPop(ctx, 30*time.Second, DocumentScraperQueue).Result()
if err != nil {
if err == redis.Nil {
return nil, nil // Queue is empty
}
q.logger.Error("Failed to dequeue job", map[string]interface{}{
"error": err.Error(),
})
return nil, err
}
if len(result) < 2 {
return nil, fmt.Errorf("invalid dequeue result")
}
jobData := result[1]
// Parse job
var job DocumentScraperJob
if err := job.FromJSON(jobData); err != nil {
q.logger.Error("Failed to parse job", map[string]interface{}{
"error": err.Error(),
"data": jobData,
})
return nil, err
}
// Mark as in progress with TTL
job.Status = string(StatusInProgress)
job.Attempts++
job.ProcessingStartedAt = time.Now().Unix()
job.UpdatedAt = time.Now().Unix()
jobKey := fmt.Sprintf("%s%s", JobStatusPrefix, job.ID)
if err := q.client.Set(ctx, jobKey, string(StatusInProgress), InProgressExpiration).Err(); err != nil {
q.logger.Warn("Failed to update job status", map[string]interface{}{
"job_id": job.ID,
"error": err.Error(),
})
}
q.logger.Info("Job dequeued successfully", map[string]interface{}{
"job_id": job.ID,
"attempt": job.Attempts,
"max_retry": job.MaxRetries,
})
return &job, nil
}
// AckJob marks a job as completed
func (q *RedisQueue) AckJob(ctx context.Context, jobID string) error {
now := time.Now().Unix()
jobKey := fmt.Sprintf("%s%s", JobStatusPrefix, jobID)
// Update status with TTL for record keeping
if err := q.client.Set(ctx, jobKey, string(StatusCompleted), 48*time.Hour).Err(); err != nil {
q.logger.Warn("Failed to ack job", map[string]interface{}{
"job_id": jobID,
"error": err.Error(),
})
}
// Store completion timestamp
completedKey := fmt.Sprintf("%s:completed:%s", DocumentScraperQueue, jobID)
if err := q.client.Set(ctx, completedKey, now, 48*time.Hour).Err(); err != nil {
q.logger.Warn("Failed to store completion timestamp", map[string]interface{}{
"job_id": jobID,
"error": err.Error(),
})
}
q.logger.Info("Job acknowledged", map[string]interface{}{
"job_id": jobID,
})
return nil
}
// NackJob marks a job as failed or for retry
func (q *RedisQueue) NackJob(ctx context.Context, jobID string, jobErr error) error {
jobKey := fmt.Sprintf("%s%s", JobStatusPrefix, jobID)
// Get the job from Redis to check retry attempts
_, err := q.client.Get(ctx, jobKey).Result()
if err != nil && err != redis.Nil {
q.logger.Warn("Failed to get job for nack", map[string]interface{}{
"job_id": jobID,
"error": err.Error(),
})
}
// Check if we should retry by trying to get the job back from processing
// For now, mark as failed
status := StatusFailed
duration := 48 * time.Hour
if err := q.client.Set(ctx, jobKey, string(status), duration).Err(); err != nil {
q.logger.Warn("Failed to update job status on nack", map[string]interface{}{
"job_id": jobID,
"error": err.Error(),
})
}
// Store error message
errorKey := fmt.Sprintf("%s:error:%s", DocumentScraperQueue, jobID)
if err := q.client.Set(ctx, errorKey, jobErr.Error(), duration).Err(); err != nil {
q.logger.Warn("Failed to store error message", map[string]interface{}{
"job_id": jobID,
"error": err.Error(),
})
}
// Optionally move to DLQ
if q.config.IsDLQEnabled {
jobData, _ := q.client.Get(ctx, JobStatusPrefix+jobID).Result()
if jobData != "" {
if err := q.client.RPush(ctx, DocumentScraperDLQ, jobData).Err(); err != nil {
q.logger.Warn("Failed to add job to DLQ", map[string]interface{}{
"job_id": jobID,
"error": err.Error(),
})
}
}
}
q.logger.Info("Job nacked", map[string]interface{}{
"job_id": jobID,
"error": jobErr.Error(),
})
return nil
}
// GetJobStatus retrieves the status of a job
func (q *RedisQueue) GetJobStatus(ctx context.Context, jobID string) (JobStatus, error) {
jobKey := fmt.Sprintf("%s%s", JobStatusPrefix, jobID)
status, err := q.client.Get(ctx, jobKey).Result()
if err != nil {
if err == redis.Nil {
return StatusPending, nil
}
return StatusPending, err
}
return JobStatus(status), nil
}
// GetMetrics returns queue metrics
func (q *RedisQueue) GetMetrics(ctx context.Context) (*QueueMetrics, error) {
metrics := &QueueMetrics{
UpdatedAt: time.Now().Unix(),
}
// Get queue size
size, err := q.GetQueueSize(ctx)
if err != nil {
q.logger.Warn("Failed to get queue size for metrics", map[string]interface{}{
"error": err.Error(),
})
}
metrics.PendingJobs = size
// Get DLQ size
dlqSize, err := q.GetDLQSize(ctx)
if err != nil {
q.logger.Warn("Failed to get DLQ size for metrics", map[string]interface{}{
"error": err.Error(),
})
}
metrics.FailedJobs = dlqSize
return metrics, nil
}
// GetQueueSize returns the number of items in the queue
func (q *RedisQueue) GetQueueSize(ctx context.Context) (int64, error) {
size, err := q.client.LLen(ctx, DocumentScraperQueue).Result()
if err != nil {
return 0, err
}
return size, nil
}
// GetDLQSize returns the number of items in the dead letter queue
func (q *RedisQueue) GetDLQSize(ctx context.Context) (int64, error) {
size, err := q.client.LLen(ctx, DocumentScraperDLQ).Result()
if err != nil {
return 0, err
}
return size, nil
}
// GetInProgressJobs returns jobs currently being processed
func (q *RedisQueue) GetInProgressJobs(ctx context.Context) ([]DocumentScraperJob, error) {
// This is a simplified implementation
// In production, you'd track this separately
return nil, nil
}
// PurgeQueue removes all jobs from the queue
func (q *RedisQueue) PurgeQueue(ctx context.Context) error {
if err := q.client.Del(ctx, DocumentScraperQueue).Err(); err != nil {
q.logger.Error("Failed to purge queue", map[string]interface{}{
"error": err.Error(),
})
return err
}
q.logger.Info("Queue purged successfully", map[string]interface{}{})
return nil
}
// PurgeDLQ removes all jobs from the dead letter queue
func (q *RedisQueue) PurgeDLQ(ctx context.Context) error {
if err := q.client.Del(ctx, DocumentScraperDLQ).Err(); err != nil {
q.logger.Error("Failed to purge DLQ", map[string]interface{}{
"error": err.Error(),
})
return err
}
q.logger.Info("DLQ purged successfully", map[string]interface{}{})
return nil
}
// Close closes the Redis connection
func (q *RedisQueue) Close() error {
return q.client.Close()
}