package queue import ( "context" "encoding/json" "fmt" "strconv" "time" "github.com/redis/go-redis/v9" ) const ( queueKey = "jobs:queue" jobKeyPrefix = "jobs:" jobTTL = 1 * time.Hour ) // JobStatus represents the state of a generation job. type JobStatus struct { JobID string `json:"job_id"` Status string `json:"status"` // pending | processing | done | error Step string `json:"step"` Message string `json:"message"` Result string `json:"result,omitempty"` Error string `json:"error,omitempty"` ElapsedMs int64 `json:"elapsed_ms,omitempty"` TimingsMs map[string]int64 `json:"timings_ms,omitempty"` } // NewRedisClient creates a Redis client and pings to verify connectivity. func NewRedisClient(addr, password string) (*redis.Client, error) { rdb := redis.NewClient(&redis.Options{ Addr: addr, Password: password, DB: 0, }) if err := rdb.Ping(context.Background()).Err(); err != nil { return nil, fmt.Errorf("redis connect: %w", err) } return rdb, nil } // EnqueueJob pushes a job ID onto the queue and sets initial status. func EnqueueJob(ctx context.Context, rdb *redis.Client, jobID string) error { pipe := rdb.Pipeline() pipe.RPush(ctx, queueKey, jobID) pipe.HSet(ctx, jobKey(jobID), "status", "pending", "queued_at", nowMillis(), "updated_at", nowMillis()) pipe.Expire(ctx, jobKey(jobID), jobTTL) _, err := pipe.Exec(ctx) return err } // DequeueJob blocks until a job is available, then pops and returns its ID. func DequeueJob(ctx context.Context, rdb *redis.Client) (string, error) { result, err := rdb.BLPop(ctx, 0, queueKey).Result() if err != nil { return "", err } if len(result) < 2 { return "", fmt.Errorf("unexpected blpop result: %v", result) } return result[1], nil } // SetJobStatus updates the job status field. func SetJobStatus(ctx context.Context, rdb *redis.Client, jobID, status string) error { fields := []any{"status", status, "updated_at", nowMillis()} if status == "processing" { fields = append(fields, "started_at", nowMillis()) } return rdb.HSet(ctx, jobKey(jobID), fields...).Err() } // SetJobProgress updates the step and message for a job. func SetJobProgress(ctx context.Context, rdb *redis.Client, jobID, step, msg string) error { return rdb.HSet(ctx, jobKey(jobID), "step", step, "message", msg, "updated_at", nowMillis()).Err() } // SetJobTiming stores a single step duration and the total elapsed time so far. func SetJobTiming(ctx context.Context, rdb *redis.Client, jobID, step string, duration, elapsed time.Duration) error { return rdb.HSet(ctx, jobKey(jobID), step+"_ms", duration.Milliseconds(), "elapsed_ms", elapsed.Milliseconds(), "updated_at", nowMillis(), ).Err() } // SetJobResult marks a job as done with the result image URL. func SetJobResult(ctx context.Context, rdb *redis.Client, jobID, url string) error { return rdb.HSet(ctx, jobKey(jobID), "status", "done", "step", "done", "message", "Complete", "result", url, "updated_at", nowMillis()).Err() } // SetJobError marks a job as errored with an error message. func SetJobError(ctx context.Context, rdb *redis.Client, jobID, errMsg string) error { return rdb.HSet(ctx, jobKey(jobID), "status", "error", "step", "error", "message", errMsg, "error", errMsg, "updated_at", nowMillis()).Err() } // GetJobStatus retrieves the full status of a job. func GetJobStatus(ctx context.Context, rdb *redis.Client, jobID string) (*JobStatus, error) { m, err := rdb.HGetAll(ctx, jobKey(jobID)).Result() if err != nil { return nil, err } if len(m) == 0 { return nil, fmt.Errorf("job not found") } js := &JobStatus{ JobID: jobID, Status: m["status"], Step: m["step"], Message: m["message"], Result: m["result"], Error: m["error"], ElapsedMs: parseInt64(m["elapsed_ms"]), TimingsMs: map[string]int64{}, } for _, step := range []string{"check", "mask", "inpaint", "cleanup", "total"} { if value := parseInt64(m[step+"_ms"]); value > 0 { js.TimingsMs[step] = value } } if len(js.TimingsMs) == 0 { js.TimingsMs = nil } return js, nil } // StoreJobPayload saves the job's input parameters and render options. func StoreJobPayload(ctx context.Context, rdb *redis.Client, jobID string, payload map[string]string) error { data, _ := json.Marshal(payload) return rdb.HSet(ctx, jobKey(jobID), "payload", string(data)).Err() } // GetJobPayload retrieves the stored input parameters. func GetJobPayload(ctx context.Context, rdb *redis.Client, jobID string) (map[string]string, error) { s, err := rdb.HGet(ctx, jobKey(jobID), "payload").Result() if err != nil { return nil, err } var p map[string]string json.Unmarshal([]byte(s), &p) return p, nil } func jobKey(id string) string { return jobKeyPrefix + id } func nowMillis() int64 { return time.Now().UnixMilli() } func parseInt64(raw string) int64 { value, _ := strconv.ParseInt(raw, 10, 64) return value }