package distributed import ( "context" "encoding/json" "fmt" "net/url" "strconv" "strings" "time" "github.com/j3ssie/osmedeus/v5/internal/config" "github.com/j3ssie/osmedeus/v5/internal/core" "github.com/redis/rueidis" ) // Redis key prefixes const ( KeyPrefix = "osm:" KeyTasksPending = KeyPrefix + "tasks:pending" KeyTasksRunning = KeyPrefix + "tasks:running" KeyTasksCompleted = KeyPrefix + "tasks:completed" KeyWorkers = KeyPrefix + "workers" KeyWorkersHeartbeat = KeyPrefix + "workers:heartbeat" KeyMasterLock = KeyPrefix + "master:lock" // Event broker keys (pub/sub channels) KeyEventsPrefix = KeyPrefix + "events:" // osm:events:{topic} // Data queue keys (for worker -> master data) KeyDataRuns = KeyPrefix + "data:runs" KeyDataSteps = KeyPrefix + "data:steps" KeyDataEvents = KeyPrefix + "data:events" KeyDataArtifacts = KeyPrefix + "data:artifacts" KeyDataExecute = KeyPrefix + "data:execute" KeyDataExecuteWorker = KeyPrefix + "data:execute:worker:" // osm:data:execute:worker:{worker_id} ) // KeyDataExecuteForWorker returns the per-worker execute queue key. func KeyDataExecuteForWorker(workerID string) string { return KeyDataExecuteWorker + workerID } // Timeouts and intervals const ( HeartbeatInterval = 30 * time.Second HeartbeatTimeout = 90 * time.Second // 3 missed heartbeats TaskPollTimeout = 5 * time.Second DefaultConnectTimeout = 60 * time.Second ) // Client wraps a rueidis client with helper methods type Client struct { client rueidis.Client cfg *config.RedisConfig } // NewClient creates a new Redis client from configuration func NewClient(cfg *config.RedisConfig) (*Client, error) { if cfg.Host == "" { return nil, fmt.Errorf("redis host not configured") } port := cfg.Port if port == 0 { port = 6379 } opts := rueidis.ClientOption{ InitAddress: []string{fmt.Sprintf("%s:%d", cfg.Host, port)}, Username: cfg.Username, Password: cfg.Password, SelectDB: cfg.DB, DisableCache: true, // Disable client-side caching for simpler behavior } client, err := rueidis.NewClient(opts) if err != nil { return nil, fmt.Errorf("failed to create redis client: %w", err) } return &Client{ client: client, cfg: cfg, }, nil } // NewClientFromConfig creates a client from the global config func NewClientFromConfig(cfg *config.Config) (*Client, error) { return NewClient(&cfg.Redis) } // ParseRedisURL parses a Redis connection URL into RedisConfig // Format: redis://[username:password@]host:port[/db] func ParseRedisURL(redisURL string) (*config.RedisConfig, error) { if !strings.HasPrefix(redisURL, "redis://") { redisURL = "redis://" + redisURL } u, err := url.Parse(redisURL) if err != nil { return nil, fmt.Errorf("invalid redis URL: %w", err) } cfg := &config.RedisConfig{ Host: u.Hostname(), Port: 6379, ConnectionTimeout: 60, } if u.Port() != "" { port, err := strconv.Atoi(u.Port()) if err != nil { return nil, fmt.Errorf("invalid redis port: %w", err) } cfg.Port = port } if u.User != nil { cfg.Username = u.User.Username() cfg.Password, _ = u.User.Password() } if u.Path != "" && u.Path != "/" { db, err := strconv.Atoi(strings.TrimPrefix(u.Path, "/")) if err == nil { cfg.DB = db } } return cfg, nil } // Close closes the Redis client func (c *Client) Close() { c.client.Close() } // Ping tests the Redis connection func (c *Client) Ping(ctx context.Context) error { cmd := c.client.B().Ping().Build() return c.client.Do(ctx, cmd).Error() } // Raw returns the underlying rueidis client func (c *Client) Raw() rueidis.Client { return c.client } // PushTask pushes a task to the pending queue func (c *Client) PushTask(ctx context.Context, task *Task) error { data, err := task.MarshalJSON() if err != nil { return fmt.Errorf("failed to marshal task: %w", err) } cmd := c.client.B().Lpush().Key(KeyTasksPending).Element(string(data)).Build() return c.client.Do(ctx, cmd).Error() } // PopTask pops a task from the pending queue (blocking) func (c *Client) PopTask(ctx context.Context, timeout time.Duration) (*Task, error) { cmd := c.client.B().Brpop().Key(KeyTasksPending).Timeout(timeout.Seconds()).Build() result, err := c.client.Do(ctx, cmd).AsStrSlice() if err != nil { if rueidis.IsRedisNil(err) { return nil, nil // Timeout, no task available } return nil, fmt.Errorf("failed to pop task: %w", err) } if len(result) < 2 { return nil, nil // No task } return UnmarshalTask([]byte(result[1])) } // SetTaskRunning moves a task to the running hash func (c *Client) SetTaskRunning(ctx context.Context, task *Task) error { data, err := task.MarshalJSON() if err != nil { return fmt.Errorf("failed to marshal task: %w", err) } cmd := c.client.B().Hset().Key(KeyTasksRunning).FieldValue().FieldValue(task.ID, string(data)).Build() return c.client.Do(ctx, cmd).Error() } // RemoveTaskRunning removes a task from the running hash func (c *Client) RemoveTaskRunning(ctx context.Context, taskID string) error { cmd := c.client.B().Hdel().Key(KeyTasksRunning).Field(taskID).Build() return c.client.Do(ctx, cmd).Error() } // SetTaskResult stores a task result in the completed hash func (c *Client) SetTaskResult(ctx context.Context, result *TaskResult) error { data, err := result.MarshalJSON() if err != nil { return fmt.Errorf("failed to marshal result: %w", err) } cmd := c.client.B().Hset().Key(KeyTasksCompleted).FieldValue().FieldValue(result.TaskID, string(data)).Build() return c.client.Do(ctx, cmd).Error() } // GetTaskResult retrieves a task result from the completed hash func (c *Client) GetTaskResult(ctx context.Context, taskID string) (*TaskResult, error) { cmd := c.client.B().Hget().Key(KeyTasksCompleted).Field(taskID).Build() data, err := c.client.Do(ctx, cmd).ToString() if err != nil { if rueidis.IsRedisNil(err) { return nil, nil } return nil, fmt.Errorf("failed to get task result: %w", err) } return UnmarshalTaskResult([]byte(data)) } // GetRunningTask retrieves a running task by ID func (c *Client) GetRunningTask(ctx context.Context, taskID string) (*Task, error) { cmd := c.client.B().Hget().Key(KeyTasksRunning).Field(taskID).Build() data, err := c.client.Do(ctx, cmd).ToString() if err != nil { if rueidis.IsRedisNil(err) { return nil, nil } return nil, fmt.Errorf("failed to get running task: %w", err) } return UnmarshalTask([]byte(data)) } // GetAllRunningTasks retrieves all running tasks func (c *Client) GetAllRunningTasks(ctx context.Context) ([]*Task, error) { cmd := c.client.B().Hgetall().Key(KeyTasksRunning).Build() result, err := c.client.Do(ctx, cmd).AsStrMap() if err != nil { return nil, fmt.Errorf("failed to get running tasks: %w", err) } var tasks []*Task for _, data := range result { task, err := UnmarshalTask([]byte(data)) if err != nil { continue } tasks = append(tasks, task) } return tasks, nil } // RegisterWorker registers a worker in the workers hash func (c *Client) RegisterWorker(ctx context.Context, worker *WorkerInfo) error { data, err := worker.MarshalJSON() if err != nil { return fmt.Errorf("failed to marshal worker: %w", err) } cmd := c.client.B().Hset().Key(KeyWorkers).FieldValue().FieldValue(worker.ID, string(data)).Build() return c.client.Do(ctx, cmd).Error() } // UpdateWorkerHeartbeat updates a worker's heartbeat timestamp func (c *Client) UpdateWorkerHeartbeat(ctx context.Context, workerID string) error { timestamp := strconv.FormatInt(time.Now().Unix(), 10) cmd := c.client.B().Hset().Key(KeyWorkersHeartbeat).FieldValue().FieldValue(workerID, timestamp).Build() return c.client.Do(ctx, cmd).Error() } // GetWorkerHeartbeat gets a worker's last heartbeat timestamp func (c *Client) GetWorkerHeartbeat(ctx context.Context, workerID string) (time.Time, error) { cmd := c.client.B().Hget().Key(KeyWorkersHeartbeat).Field(workerID).Build() data, err := c.client.Do(ctx, cmd).ToString() if err != nil { if rueidis.IsRedisNil(err) { return time.Time{}, nil } return time.Time{}, err } ts, err := strconv.ParseInt(data, 10, 64) if err != nil { return time.Time{}, err } return time.Unix(ts, 0), nil } // GetAllWorkers retrieves all registered workers func (c *Client) GetAllWorkers(ctx context.Context) ([]*WorkerInfo, error) { cmd := c.client.B().Hgetall().Key(KeyWorkers).Build() result, err := c.client.Do(ctx, cmd).AsStrMap() if err != nil { return nil, fmt.Errorf("failed to get workers: %w", err) } var workers []*WorkerInfo for _, data := range result { worker, err := UnmarshalWorkerInfo([]byte(data)) if err != nil { continue } workers = append(workers, worker) } return workers, nil } // GetWorker retrieves a single worker by ID func (c *Client) GetWorker(ctx context.Context, workerID string) (*WorkerInfo, error) { cmd := c.client.B().Hget().Key(KeyWorkers).Field(workerID).Build() data, err := c.client.Do(ctx, cmd).ToString() if err != nil { if rueidis.IsRedisNil(err) { return nil, nil } return nil, fmt.Errorf("failed to get worker: %w", err) } return UnmarshalWorkerInfo([]byte(data)) } // GetWorkerByAlias retrieves a worker by its alias func (c *Client) GetWorkerByAlias(ctx context.Context, alias string) (*WorkerInfo, error) { workers, err := c.GetAllWorkers(ctx) if err != nil { return nil, err } for _, w := range workers { if w.Alias == alias { return w, nil } } return nil, nil } // RemoveWorker removes a worker from the registry func (c *Client) RemoveWorker(ctx context.Context, workerID string) error { // Remove from both workers and heartbeat hashes cmd1 := c.client.B().Hdel().Key(KeyWorkers).Field(workerID).Build() cmd2 := c.client.B().Hdel().Key(KeyWorkersHeartbeat).Field(workerID).Build() if err := c.client.Do(ctx, cmd1).Error(); err != nil { return err } return c.client.Do(ctx, cmd2).Error() } // AcquireMasterLock tries to acquire the master lock func (c *Client) AcquireMasterLock(ctx context.Context, masterID string, ttl time.Duration) (bool, error) { cmd := c.client.B().Set().Key(KeyMasterLock).Value(masterID).Nx().Ex(ttl).Build() result, err := c.client.Do(ctx, cmd).ToString() if err != nil { if rueidis.IsRedisNil(err) { return false, nil // Lock not acquired } return false, err } return result == "OK", nil } // RefreshMasterLock refreshes the master lock TTL func (c *Client) RefreshMasterLock(ctx context.Context, masterID string, ttl time.Duration) error { // Only refresh if we still own the lock cmd := c.client.B().Get().Key(KeyMasterLock).Build() current, err := c.client.Do(ctx, cmd).ToString() if err != nil { return err } if current != masterID { return fmt.Errorf("master lock lost") } expireCmd := c.client.B().Expire().Key(KeyMasterLock).Seconds(int64(ttl.Seconds())).Build() return c.client.Do(ctx, expireCmd).Error() } // ReleaseMasterLock releases the master lock func (c *Client) ReleaseMasterLock(ctx context.Context, masterID string) error { // Only release if we own the lock cmd := c.client.B().Get().Key(KeyMasterLock).Build() current, err := c.client.Do(ctx, cmd).ToString() if err != nil { if rueidis.IsRedisNil(err) { return nil // Already released } return err } if current != masterID { return nil // Not our lock } delCmd := c.client.B().Del().Key(KeyMasterLock).Build() return c.client.Do(ctx, delCmd).Error() } // ============================================================================= // Event Pub/Sub Methods // ============================================================================= // PublishEvent publishes an event to a topic channel func (c *Client) PublishEvent(ctx context.Context, topic string, event *core.Event) error { data, err := json.Marshal(event) if err != nil { return fmt.Errorf("failed to marshal event: %w", err) } channel := KeyEventsPrefix + topic cmd := c.client.B().Publish().Channel(channel).Message(string(data)).Build() return c.client.Do(ctx, cmd).Error() } // SubscribeEvents subscribes to event channels with pattern and calls handler for each event. // This method blocks until the context is cancelled or an error occurs. func (c *Client) SubscribeEvents(ctx context.Context, handler func(*core.Event)) error { pattern := KeyEventsPrefix + "*" // Use Receive with PSUBSCRIBE - this blocks and calls handler for each message err := c.client.Receive(ctx, c.client.B().Psubscribe().Pattern(pattern).Build(), func(msg rueidis.PubSubMessage) { if msg.Message != "" { var event core.Event if err := json.Unmarshal([]byte(msg.Message), &event); err == nil { handler(&event) } } }) if err != nil && ctx.Err() != nil { return ctx.Err() } return err } // ============================================================================= // Data Queue Methods (Worker -> Master) // ============================================================================= // DataEnvelope wraps data with type information for queue processing type DataEnvelope struct { Type string `json:"type"` Data json.RawMessage `json:"data"` Timestamp time.Time `json:"timestamp"` WorkerID string `json:"worker_id,omitempty"` } // PushData pushes data to a worker data queue func (c *Client) PushData(ctx context.Context, key string, dataType string, data interface{}, workerID string) error { dataBytes, err := json.Marshal(data) if err != nil { return fmt.Errorf("failed to marshal data: %w", err) } envelope := DataEnvelope{ Type: dataType, Data: dataBytes, Timestamp: time.Now(), WorkerID: workerID, } envBytes, err := json.Marshal(envelope) if err != nil { return fmt.Errorf("failed to marshal envelope: %w", err) } cmd := c.client.B().Lpush().Key(key).Element(string(envBytes)).Build() return c.client.Do(ctx, cmd).Error() } // PopData pops data from a queue (blocking with timeout) func (c *Client) PopData(ctx context.Context, key string, timeout time.Duration) (*DataEnvelope, error) { cmd := c.client.B().Brpop().Key(key).Timeout(timeout.Seconds()).Build() result, err := c.client.Do(ctx, cmd).AsStrSlice() if err != nil { if rueidis.IsRedisNil(err) { return nil, nil // Timeout, no data available } return nil, fmt.Errorf("failed to pop data: %w", err) } if len(result) < 2 { return nil, nil // No data } var envelope DataEnvelope if err := json.Unmarshal([]byte(result[1]), &envelope); err != nil { return nil, fmt.Errorf("failed to unmarshal envelope: %w", err) } return &envelope, nil } // PopDataMulti pops data from multiple queues (blocking with timeout) // Returns the key that had data and the data envelope func (c *Client) PopDataMulti(ctx context.Context, timeout time.Duration, keys ...string) (string, *DataEnvelope, error) { cmd := c.client.B().Brpop().Key(keys[0]) for _, k := range keys[1:] { cmd = cmd.Key(k) } cmdBuilt := cmd.Timeout(timeout.Seconds()).Build() result, err := c.client.Do(ctx, cmdBuilt).AsStrSlice() if err != nil { if rueidis.IsRedisNil(err) { return "", nil, nil // Timeout, no data available } return "", nil, fmt.Errorf("failed to pop data: %w", err) } if len(result) < 2 { return "", nil, nil // No data } var envelope DataEnvelope if err := json.Unmarshal([]byte(result[1]), &envelope); err != nil { return "", nil, fmt.Errorf("failed to unmarshal envelope: %w", err) } return result[0], &envelope, nil } // GetQueueLength returns the length of a data queue func (c *Client) GetQueueLength(ctx context.Context, key string) (int64, error) { cmd := c.client.B().Llen().Key(key).Build() return c.client.Do(ctx, cmd).AsInt64() }