agent SDK fix
This commit is contained in:
m.nazemi
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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"`
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
Reference in New Issue
Block a user