...

pkg/servers/heroagent/redis.go

@@ -0,0 +1,199 @@
+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>)
+		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)
+}