package distributed import ( "context" "fmt" "os" "os/exec" "strings" "sync" "time" "github.com/j3ssie/osmedeus/v5/internal/json" "github.com/google/uuid" "github.com/j3ssie/osmedeus/v5/internal/broker" "github.com/j3ssie/osmedeus/v5/internal/config" "github.com/j3ssie/osmedeus/v5/internal/core" "github.com/j3ssie/osmedeus/v5/internal/database" "github.com/j3ssie/osmedeus/v5/internal/database/repository" "github.com/j3ssie/osmedeus/v5/internal/executor" "github.com/j3ssie/osmedeus/v5/internal/functions" "github.com/j3ssie/osmedeus/v5/internal/logger" "github.com/j3ssie/osmedeus/v5/internal/terminal" "github.com/uptrace/bun" "go.uber.org/zap" ) const ( MasterLockTTL = 60 * time.Second MasterLockRefresh = 30 * time.Second WorkerCheckPeriod = 30 * time.Second ) // EventHandler is called when an event is received via Redis subscription type EventHandler func(*core.Event) // Master represents a master node that coordinates workers type Master struct { ID string client *Client config *config.Config logger *zap.Logger printer *terminal.Printer // Event handling eventHandler EventHandler eventBroker *broker.RedisEventBroker // Database for persisting worker data db *bun.DB // For tracking mu sync.RWMutex running bool } // NewMaster creates a new master node func NewMaster(cfg *config.Config) (*Master, error) { client, err := NewClientFromConfig(cfg) if err != nil { return nil, fmt.Errorf("failed to create redis client: %w", err) } hostname, _ := os.Hostname() masterID := fmt.Sprintf("master-%s-%s", hostname, uuid.NewString()[:8]) log := logger.Get() // Initialize event broker eventBroker, err := broker.GetSharedBroker() if err != nil { log.Warn("Failed to initialize event broker", zap.Error(err)) } return &Master{ ID: masterID, client: client, config: cfg, logger: log, printer: terminal.NewPrinter(), eventBroker: eventBroker, db: database.GetDB(), }, nil } // SetEventHandler sets the callback function for handling received events func (m *Master) SetEventHandler(handler EventHandler) { m.mu.Lock() defer m.mu.Unlock() m.eventHandler = handler } // Start starts the master node func (m *Master) Start(ctx context.Context) error { // Test connection if err := m.client.Ping(ctx); err != nil { return fmt.Errorf("failed to connect to redis: %w", err) } // Acquire master lock acquired, err := m.client.AcquireMasterLock(ctx, m.ID, MasterLockTTL) if err != nil { return fmt.Errorf("failed to acquire master lock: %w", err) } if !acquired { return fmt.Errorf("another master is already running") } m.mu.Lock() m.running = true m.mu.Unlock() m.printer.Success("Master %s started", m.ID) m.printer.Info("Redis connected at %s", terminal.Cyan(fmt.Sprintf("%s:%d", m.config.Redis.Host, m.config.Redis.Port))) m.printer.Info("Waiting for workers and tasks...") // Start lock refresh goroutine lockCtx, cancelLock := context.WithCancel(ctx) defer cancelLock() go m.lockRefreshLoop(lockCtx) // Start worker monitor goroutine monitorCtx, cancelMonitor := context.WithCancel(ctx) defer cancelMonitor() go m.workerMonitorLoop(monitorCtx) // Start event subscription loop (for Redis pub/sub events) if m.eventBroker != nil { eventCtx, cancelEvent := context.WithCancel(ctx) defer cancelEvent() go m.eventSubscriptionLoop(eventCtx) m.printer.Info("Event subscription started") } // Start data processor loop (for worker data queues) dataCtx, cancelData := context.WithCancel(ctx) defer cancelData() go m.dataProcessorLoop(dataCtx) m.printer.Info("Data processor started") // Wait for shutdown <-ctx.Done() m.printer.Info("Master %s shutting down...", terminal.Cyan(m.ID)) m.cleanup(context.Background()) return nil } // lockRefreshLoop periodically refreshes the master lock func (m *Master) lockRefreshLoop(ctx context.Context) { ticker := time.NewTicker(MasterLockRefresh) defer ticker.Stop() for { select { case <-ctx.Done(): return case <-ticker.C: if err := m.client.RefreshMasterLock(ctx, m.ID, MasterLockTTL); err != nil { m.printer.Error("Failed to refresh master lock: %s", err) // If we lose the lock, we should stop m.mu.Lock() m.running = false m.mu.Unlock() return } } } } // workerMonitorLoop monitors worker heartbeats and handles failures func (m *Master) workerMonitorLoop(ctx context.Context) { ticker := time.NewTicker(WorkerCheckPeriod) defer ticker.Stop() for { select { case <-ctx.Done(): return case <-ticker.C: m.checkWorkerHealth(ctx) } } } // checkWorkerHealth checks for dead workers and reassigns their tasks func (m *Master) checkWorkerHealth(ctx context.Context) { workers, err := m.client.GetAllWorkers(ctx) if err != nil { m.printer.Warning("Failed to get workers: %s", err) return } now := time.Now() for _, worker := range workers { heartbeat, err := m.client.GetWorkerHeartbeat(ctx, worker.ID) if err != nil { continue } // Check if worker is dead (missed heartbeats) if heartbeat.IsZero() || now.Sub(heartbeat) > HeartbeatTimeout { m.printer.Warning("Worker %s appears dead (last heartbeat %s ago), reassigning tasks...", terminal.Cyan(worker.ID), now.Sub(heartbeat).Round(time.Second)) // Reassign the worker's running tasks m.reassignWorkerTasks(ctx, worker.ID) // Remove the dead worker if err := m.client.RemoveWorker(ctx, worker.ID); err != nil { m.printer.Warning("Failed to remove dead worker: %s", err) } } } } // reassignWorkerTasks moves a dead worker's tasks back to pending func (m *Master) reassignWorkerTasks(ctx context.Context, workerID string) { tasks, err := m.client.GetAllRunningTasks(ctx) if err != nil { m.printer.Error("Failed to get running tasks: %s", err) return } for _, task := range tasks { if task.WorkerID == workerID { m.printer.Info("Reassigning task %s from worker %s", terminal.Cyan(task.ID), terminal.Cyan(workerID)) // Reset task status task.Status = TaskStatusPending task.WorkerID = "" task.StartedAt = nil // Push back to pending queue if err := m.client.PushTask(ctx, task); err != nil { m.printer.Error("Failed to reassign task %s: %s", task.ID, err) continue } // Remove from running if err := m.client.RemoveTaskRunning(ctx, task.ID); err != nil { m.printer.Warning("Failed to remove task from running: %s", err) } } } } // cleanup releases the master lock func (m *Master) cleanup(ctx context.Context) { m.printer.Info("Cleaning up master %s...", m.ID) m.mu.Lock() m.running = false m.mu.Unlock() if err := m.client.ReleaseMasterLock(ctx, m.ID); err != nil { m.printer.Warning("Failed to release master lock: %s", err) } m.client.Close() } // SubmitTask submits a new task to the pending queue func (m *Master) SubmitTask(ctx context.Context, task *Task) error { if task.ID == "" { task.ID = uuid.NewString()[:8] } if task.CreatedAt.IsZero() { task.CreatedAt = time.Now() } task.Status = TaskStatusPending m.printer.Info("Submitting task %s: %s -> %s", terminal.Cyan(task.ID), terminal.Yellow(task.WorkflowName), terminal.Green(task.Target)) return m.client.PushTask(ctx, task) } // GetTaskStatus retrieves the status of a task func (m *Master) GetTaskStatus(ctx context.Context, taskID string) (*Task, *TaskResult, error) { // Check running tasks first task, err := m.client.GetRunningTask(ctx, taskID) if err != nil { return nil, nil, err } if task != nil { return task, nil, nil } // Check completed tasks result, err := m.client.GetTaskResult(ctx, taskID) if err != nil { return nil, nil, err } if result != nil { return nil, result, nil } return nil, nil, fmt.Errorf("task not found: %s", taskID) } // ListWorkers returns all registered workers with their current status func (m *Master) ListWorkers(ctx context.Context) ([]*WorkerInfo, error) { workers, err := m.client.GetAllWorkers(ctx) if err != nil { return nil, err } // Enrich with heartbeat info now := time.Now() for _, worker := range workers { heartbeat, err := m.client.GetWorkerHeartbeat(ctx, worker.ID) if err == nil { worker.LastHeartbeat = heartbeat // Update status based on heartbeat if now.Sub(heartbeat) > HeartbeatTimeout { worker.Status = "offline" } } } return workers, nil } // ListTasks returns all tasks (running and completed) func (m *Master) ListTasks(ctx context.Context) ([]*Task, []*TaskResult, error) { running, err := m.client.GetAllRunningTasks(ctx) if err != nil { return nil, nil, err } // Get completed tasks (we'd need to iterate the hash) // For now, return running tasks return running, nil, nil } // GetClient returns the Redis client for external use func (m *Master) GetClient() *Client { return m.client } // IsRunning returns whether the master is currently running func (m *Master) IsRunning() bool { m.mu.RLock() defer m.mu.RUnlock() return m.running } // ============================================================================= // Event Subscription Loop // ============================================================================= // eventSubscriptionLoop subscribes to Redis pub/sub events and forwards them to the handler func (m *Master) eventSubscriptionLoop(ctx context.Context) { err := m.eventBroker.SubscribeEvents(ctx, func(event *core.Event) { m.logger.Debug("received event via Redis", zap.String("topic", event.Topic), zap.String("source", event.Source), zap.String("event_id", event.ID), ) // Forward to registered handler (e.g., EventReceiver) m.mu.RLock() handler := m.eventHandler m.mu.RUnlock() if handler != nil { handler(event) } // Also persist to database m.persistEventLog(ctx, event) }) if err != nil && ctx.Err() == nil { m.printer.Error("Event subscription error: %s", err) } } // persistEventLog saves an event to the database func (m *Master) persistEventLog(ctx context.Context, event *core.Event) { if m.db == nil { return } eventLog := &database.EventLog{ Topic: event.Topic, EventID: event.ID, Name: event.Name, Source: event.Source, DataType: event.DataType, Data: event.Data, Processed: true, // Events from Redis pub/sub are processed immediately CreatedAt: event.Timestamp, } repo := repository.NewEventLogRepository(m.db) if err := repo.Create(ctx, eventLog); err != nil { m.logger.Debug("failed to persist event log", zap.Error(err)) } } // ============================================================================= // Data Processor Loop // ============================================================================= // dataProcessorLoop processes data from worker data queues func (m *Master) dataProcessorLoop(ctx context.Context) { keys := []string{KeyDataRuns, KeyDataSteps, KeyDataEvents, KeyDataArtifacts, KeyDataExecute} timeout := 1 * time.Second for { select { case <-ctx.Done(): return default: // Try to pop from any of the data queues key, envelope, err := m.client.PopDataMulti(ctx, timeout, keys...) if err != nil { if ctx.Err() == nil { m.printer.Warning("Error popping from data queue: %s", err) } continue } if envelope != nil { m.processWorkerData(ctx, key, envelope) } } } } // processWorkerData processes data received from a worker func (m *Master) processWorkerData(ctx context.Context, key string, envelope *DataEnvelope) { m.logger.Debug("processing worker data", zap.String("key", key), zap.String("type", envelope.Type), zap.String("worker_id", envelope.WorkerID), ) // Execute requests don't require a database connection if key == KeyDataExecute { m.processExecuteData(ctx, envelope) return } if m.db == nil { m.db = database.GetDB() } if m.db == nil { m.logger.Debug("skipping data processing - no database connection", zap.String("key", key), zap.String("type", envelope.Type), ) return } switch key { case KeyDataRuns: m.processRunData(ctx, envelope) case KeyDataSteps: m.processStepData(ctx, envelope) case KeyDataEvents: m.processEventData(ctx, envelope) case KeyDataArtifacts: m.processArtifactData(ctx, envelope) default: m.printer.Warning("Unknown data queue key: %s", key) } } // processRunData processes run data from a worker func (m *Master) processRunData(ctx context.Context, envelope *DataEnvelope) { var run database.Run if err := json.Unmarshal(envelope.Data, &run); err != nil { m.printer.Warning("Failed to unmarshal run data: %s", err) return } repo := repository.NewRunRepository(m.db) // Check if run exists (by run_uuid) existing, err := repo.GetByRunID(ctx, run.RunUUID) if err == nil && existing != nil { // Update existing run run.ID = existing.ID if err := repo.Update(ctx, &run); err != nil { m.printer.Warning("Failed to update run %s: %s", run.RunUUID, err) } } else { // Create new run if err := repo.Create(ctx, &run); err != nil { m.printer.Warning("Failed to create run %s: %s", run.RunUUID, err) } } } // processStepData processes step result data from a worker func (m *Master) processStepData(ctx context.Context, envelope *DataEnvelope) { var step database.StepResult if err := json.Unmarshal(envelope.Data, &step); err != nil { m.printer.Warning("Failed to unmarshal step data: %s", err) return } // Resolve the real DB RunID using RunUUID sent by the worker if step.RunUUID != "" { repo := repository.NewRunRepository(m.db) existing, err := repo.GetByRunID(ctx, step.RunUUID) if err == nil && existing != nil { step.RunID = existing.ID } } // Insert step result and increment progress if step.RunID > 0 { _, err := m.db.NewInsert().Model(&step).Exec(ctx) if err != nil { m.printer.Warning("Failed to create step result %s: %s", step.StepName, err) } else { // Increment completed_steps on the run _, _ = m.db.NewUpdate().Model((*database.Run)(nil)). Set("completed_steps = completed_steps + 1"). Where("id = ?", step.RunID). Exec(ctx) } } } // processEventData processes event log data from a worker func (m *Master) processEventData(ctx context.Context, envelope *DataEnvelope) { var eventLog database.EventLog if err := json.Unmarshal(envelope.Data, &eventLog); err != nil { m.printer.Warning("Failed to unmarshal event data: %s", err) return } repo := repository.NewEventLogRepository(m.db) if err := repo.Create(ctx, &eventLog); err != nil { m.printer.Warning("Failed to create event log for topic %s: %s", eventLog.Topic, err) } } // processArtifactData processes artifact data from a worker func (m *Master) processArtifactData(ctx context.Context, envelope *DataEnvelope) { var artifact database.Artifact if err := json.Unmarshal(envelope.Data, &artifact); err != nil { m.printer.Warning("Failed to unmarshal artifact data: %s", err) return } // Insert artifact _, err := m.db.NewInsert().Model(&artifact).Exec(ctx) if err != nil { m.printer.Warning("Failed to create artifact %s: %s", artifact.Name, err) } } // ============================================================================= // Execute Request Processing // ============================================================================= // processExecuteData processes an execute request from a worker func (m *Master) processExecuteData(ctx context.Context, envelope *DataEnvelope) { var req ExecuteRequest if err := json.Unmarshal(envelope.Data, &req); err != nil { m.printer.Warning("Failed to unmarshal execute request: %s", err) return } // Resolve execute type: prefer new ExecuteType field, fall back to legacy Action executeType := req.ExecuteType if executeType == "" { executeType = req.Action } // Resolve target role: default to "master" for backward compatibility targetRole := req.TargetRole if targetRole == "" { targetRole = "master" } m.logger.Info("processing execute request from worker", zap.String("worker_id", envelope.WorkerID), zap.String("execute_type", executeType), zap.String("target_role", targetRole), zap.String("data", req.Data), ) // Route requests not targeted at master to workers if targetRole != "master" { m.routeExecuteToWorkers(ctx, &req, envelope.WorkerID) return } switch executeType { case "func": // Resolve data: prefer new Data field, fall back to legacy Expr expr := req.Data if expr == "" { expr = req.Expr } m.executeFunc(ctx, expr, envelope.WorkerID) case "run": // Resolve data: prefer new Data field, fall back to legacy Workflow workflow := req.Data if workflow == "" { workflow = req.Workflow } m.executeRun(ctx, workflow, req.Target, req.Params, envelope.WorkerID) case "bash": command := req.Data if command == "" { command = req.Expr // legacy fallback } m.executeBash(ctx, command, envelope.WorkerID) case "sync_to_worker": // Data=src path on master, Target=dest path on worker m.executeSyncToWorker(ctx, req.Data, req.Target, envelope.WorkerID) default: m.printer.Warning("Unknown execute type from worker %s: %s", envelope.WorkerID, executeType) } } // executeFunc runs a utility function expression on the master func (m *Master) executeFunc(ctx context.Context, expr, workerID string) { if expr == "" { m.printer.Warning("Empty expression from worker %s", workerID) return } m.logger.Debug("executing function from worker", zap.String("worker_id", workerID), zap.String("expr", expr), ) // Build context with built-in variables so functions have access to paths execCtx := executor.BuildBuiltinVariables(m.config, nil) registry := functions.NewRegistry() _, err := registry.Execute(expr, execCtx) if err != nil { m.printer.Warning("Failed to execute function from worker %s: %s (expr: %s)", workerID, err, expr) return } m.logger.Info("executed function from worker", zap.String("worker_id", workerID), zap.String("expr", expr), ) } // executeRun submits a workflow task to the pending queue on behalf of a worker func (m *Master) executeRun(ctx context.Context, workflow, target, params, workerID string) { if workflow == "" || target == "" { m.printer.Warning("Missing workflow or target from worker %s", workerID) return } m.logger.Debug("submitting run request from worker", zap.String("worker_id", workerID), zap.String("workflow", workflow), zap.String("target", target), zap.String("params", params), ) // Parse params into map taskParams := make(map[string]interface{}) if params != "" { for _, pair := range strings.Split(params, ",") { pair = strings.TrimSpace(pair) if pair == "" { continue } parts := strings.SplitN(pair, "=", 2) if len(parts) == 2 { taskParams[strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1]) } } } task := NewTask( uuid.NewString()[:8], workflow, "module", // default to module; master can detect flow if needed target, taskParams, ) if err := m.SubmitTask(ctx, task); err != nil { m.printer.Warning("Failed to submit task from worker %s: %s", workerID, err) return } m.logger.Info("submitted task from worker", zap.String("worker_id", workerID), zap.String("task_id", task.ID), zap.String("workflow", workflow), zap.String("target", target), ) } // executeBash runs a shell command on the master on behalf of a worker func (m *Master) executeBash(ctx context.Context, command, workerID string) { if command == "" { m.printer.Warning("Empty bash command from worker %s", workerID) return } m.logger.Debug("executing bash from worker", zap.String("worker_id", workerID), zap.String("command", command)) // @NOTE: This is intentional - bash commands come from trusted workflow YAML files // via the distributed system. Workers send requests through run_on_master('bash', cmd). cmd := exec.CommandContext(ctx, "sh", "-c", command) output, err := cmd.CombinedOutput() if err != nil { m.printer.Warning("Bash execution failed from worker %s: %s (output: %s)", workerID, err, string(output)) return } m.logger.Info("executed bash from worker", zap.String("worker_id", workerID), zap.String("command", command), zap.Int("output_len", len(output))) } // executeSyncToWorker rsyncs a file/folder from the master to the requesting worker. // src is the path on the master, dest is the path on the worker. func (m *Master) executeSyncToWorker(ctx context.Context, src, dest, workerID string) { if src == "" || dest == "" { m.printer.Warning("sync_to_worker: missing src or dest from worker %s", workerID) return } // Look up the requesting worker's SSH info worker, err := m.client.GetWorker(ctx, workerID) if err != nil || worker == nil { m.printer.Warning("sync_to_worker: failed to get worker %s: %v", workerID, err) return } if !worker.SSHEnabled { m.printer.Warning("sync_to_worker: worker %s does not have SSH enabled", workerID) return } host := worker.PublicIP if host == "" { host = worker.IPAddress } if host == "" { m.printer.Warning("sync_to_worker: worker %s has no IP address", workerID) return } m.logger.Debug("syncing to worker", zap.String("worker_id", workerID), zap.String("host", host), zap.String("src", src), zap.String("dest", dest), ) // Build rsync command: rsync -avz -e "ssh -i -p 22 -o StrictHostKeyChecking=no" src root@host:dest args := []string{"-avz", "-e"} keyPath := worker.SSHKeysPath if keyPath != "" { args = append(args, fmt.Sprintf("ssh -i %s -p 22 -o StrictHostKeyChecking=no -o ConnectTimeout=30", keyPath)) } else { args = append(args, "ssh -p 22 -o StrictHostKeyChecking=no -o ConnectTimeout=30") } rsyncDest := fmt.Sprintf("root@%s:%s", host, dest) args = append(args, src, rsyncDest) rsyncCtx, cancel := context.WithTimeout(ctx, 5*time.Minute) defer cancel() cmd := exec.CommandContext(rsyncCtx, "rsync", args...) output, err := cmd.CombinedOutput() if err != nil { m.printer.Warning("sync_to_worker: rsync to worker %s failed: %s (output: %s)", workerID, err, string(output)) return } m.logger.Info("synced to worker", zap.String("worker_id", workerID), zap.String("host", host), zap.String("src", src), zap.String("dest", dest), ) } // routeExecuteToWorkers routes an execute request to target worker(s) via per-worker queues. func (m *Master) routeExecuteToWorkers(ctx context.Context, req *ExecuteRequest, senderWorkerID string) { workers, err := m.client.GetAllWorkers(ctx) if err != nil { m.printer.Warning("Failed to get workers for routing: %s", err) return } scope := req.TargetScope if scope == "" { scope = "all" } var targets []*WorkerInfo if scope == "all" { targets = workers } else { for _, w := range workers { if w.ID == scope || w.Alias == scope || w.PublicIP == scope { targets = append(targets, w) break } } } if len(targets) == 0 { m.printer.Warning("No workers matched scope %q for execute request from %s", scope, senderWorkerID) return } for _, w := range targets { key := KeyDataExecuteForWorker(w.ID) if err := m.client.PushData(ctx, key, "execute", req, senderWorkerID); err != nil { m.printer.Warning("Failed to route execute to worker %s: %s", w.ID, err) } } m.logger.Info("routed execute request to workers", zap.String("scope", scope), zap.Int("target_count", len(targets)), zap.String("sender", senderWorkerID)) }