package heroagent import ( "context" "encoding/json" "fmt" "log" "strconv" "time" "github.com/redis/go-redis/v9" ) // RedisJobManager handles Redis operations for jobs type RedisJobManager struct { client *redis.Client ctx context.Context } // NewRedisJobManager creates a new Redis job manager func NewRedisJobManager(tcpPort int, unixSocketPath string) (*RedisJobManager, error) { // Determine network type and address var networkType, addr string if unixSocketPath != "" { networkType = "unix" addr = unixSocketPath } else { networkType = "tcp" addr = fmt.Sprintf("localhost:%d", tcpPort) } // Create Redis client client := redis.NewClient(&redis.Options{ Network: networkType, Addr: addr, DB: 0, DialTimeout: 5 * time.Second, ReadTimeout: 5 * time.Second, WriteTimeout: 5 * time.Second, }) // Test connection ctx := context.Background() _, err := client.Ping(ctx).Result() if err != nil { return nil, fmt.Errorf("failed to connect to Redis: %w", err) } return &RedisJobManager{ client: client, ctx: ctx, }, nil } // Close closes the Redis client func (r *RedisJobManager) Close() error { return r.client.Close() } // QueueKey returns the Redis queue key for a topic func QueueKey(topic string) string { return fmt.Sprintf("heroqueue:%s", topic) } // StorageKey returns the Redis storage key for a job func StorageKey(jobID uint32, topic string) string { return fmt.Sprintf("herojobs:%s:%d", topic, jobID) } // StoreJob stores a job in Redis func (r *RedisJobManager) StoreJob(job *Job) error { // Convert job to JSON jobJSON, err := json.Marshal(job) if err != nil { return fmt.Errorf("failed to marshal job: %w", err) } // Store job in Redis storageKey := StorageKey(job.JobID, job.Topic) err = r.client.Set(r.ctx, storageKey, jobJSON, 0).Err() if err != nil { return fmt.Errorf("failed to store job in Redis: %w", err) } return nil } // EnqueueJob adds a job to its queue func (r *RedisJobManager) EnqueueJob(job *Job) error { // Store the job first if err := r.StoreJob(job); err != nil { return err } // Add job ID to queue queueKey := QueueKey(job.Topic) err := r.client.RPush(r.ctx, queueKey, job.JobID).Err() if err != nil { return fmt.Errorf("failed to enqueue job: %w", err) } log.Printf("Job %d enqueued in Redis queue %s", job.JobID, queueKey) return nil } // GetJob retrieves a job from Redis func (r *RedisJobManager) GetJob(jobID uint32, topic string) (*Job, error) { // Get job from Redis storageKey := StorageKey(jobID, topic) jobJSON, err := r.client.Get(r.ctx, storageKey).Result() if err != nil { if err == redis.Nil { return nil, fmt.Errorf("job not found: %d", jobID) } return nil, fmt.Errorf("failed to get job from Redis: %w", err) } // Parse job JSON job := &Job{} if err := json.Unmarshal([]byte(jobJSON), job); err != nil { return nil, fmt.Errorf("failed to unmarshal job: %w", err) } return job, nil } // DeleteJob deletes a job from Redis func (r *RedisJobManager) DeleteJob(jobID uint32, topic string) error { // Delete job from Redis storageKey := StorageKey(jobID, topic) err := r.client.Del(r.ctx, storageKey).Err() if err != nil { return fmt.Errorf("failed to delete job from Redis: %w", err) } log.Printf("Job %d deleted from Redis", jobID) return nil } // FetchNextJob fetches the next job from a queue func (r *RedisJobManager) FetchNextJob(topic string) (*Job, error) { queueKey := QueueKey(topic) // Get and remove first job ID from queue jobIDStr, err := r.client.LPop(r.ctx, queueKey).Result() if err != nil { if err == redis.Nil { return nil, fmt.Errorf("queue is empty") } return nil, fmt.Errorf("failed to fetch job ID from queue: %w", err) } // Convert job ID to uint32 jobID, err := strconv.ParseUint(jobIDStr, 10, 32) if err != nil { return nil, fmt.Errorf("invalid job ID: %s", jobIDStr) } // Get job from Redis return r.GetJob(uint32(jobID), topic) } // ListQueues lists all job queues func (r *RedisJobManager) ListQueues() ([]string, error) { // Get all queue keys queueKeys, err := r.client.Keys(r.ctx, "heroqueue:*").Result() if err != nil { return nil, fmt.Errorf("failed to list queues: %w", err) } // Extract topic names from queue keys topics := make([]string, 0, len(queueKeys)) for _, queueKey := range queueKeys { // Extract topic from queue key (format: heroqueue:) topic := queueKey[10:] // Skip "heroqueue:" topics = append(topics, topic) } return topics, nil } // QueueSize returns the size of a queue func (r *RedisJobManager) QueueSize(topic string) (int64, error) { queueKey := QueueKey(topic) // Get queue size size, err := r.client.LLen(r.ctx, queueKey).Result() if err != nil { return 0, fmt.Errorf("failed to get queue size: %w", err) } return size, nil } // UpdateJobStatus updates the status of a job in Redis func (r *RedisJobManager) UpdateJobStatus(job *Job) error { // Update job in Redis return r.StoreJob(job) }