Files
tm_back/pkg/queue/queue.go
T
2026-04-12 12:29:35 +03:30

358 lines
9.7 KiB
Go

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()
}