Files
osmedeus/internal/distributed/master.go
T
Shai Rod df8ad1299b 🐛 fix(distributed): link worker step results to master run via UUID
Workers now propagate RunUUID from the submitted task so the master
can resolve the correct DB run record and associate step results.
Also increments completed_steps on successful insert and connects
the worker to the database for db_import_* functions.
2026-05-07 20:29:26 +03:00

843 lines
23 KiB
Go

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 <key> -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))
}