flatrender/services/node-agent/cmd/agent/main.go

// FlatRender V2 Node Agent
//
// Runs on every Windows render node. Registers itself with the V2 render
// orchestrator, sends heartbeats, claims render jobs, executes them via
// After Effects (or mock), and reports progress / completion back.
//
// Required environment variables:
//   NODE_ID             — UUID of this node (pre-created in render.render_nodes)
//
// Optional environment variables (all have sensible defaults):
//   ORCHESTRATOR_URL    — gateway base URL (default: http://localhost:8088)
//   NODE_HMAC_SECRET    — shared secret (default: node-secret-change-me)
//   NODE_REGION         — region label, e.g. "iran-tehran-1"
//   AE_PATH             — path to aerender.exe; empty = mock render
//   WORK_DIR            — scratch directory (default: system temp)
//   AGENT_VERSION       — semver string (default: 0.1.0)
//   AE_VERSION          — AE version string reported to orchestrator (default: 2024)
//   HEARTBEAT_INTERVAL_SEC — seconds between heartbeats (default: 5)
//   POLL_INTERVAL_SEC   — seconds between job-claim attempts when idle (default: 3)
//   LISTEN_PORT         — port for the health endpoint (default: 7777)
package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"sync"
	"syscall"
	"time"

	"github.com/flatrender/node-agent/internal/client"
	"github.com/flatrender/node-agent/internal/config"
	"github.com/flatrender/node-agent/internal/runner"
)

// ── Agent state ───────────────────────────────────────────────────────────────

type Agent struct {
	cfg        *config.Config
	orch       *client.Client
	mu         sync.Mutex
	currentJob *client.ClaimedJob
	status     string // "Ready" | "Busy"
}

func newAgent(cfg *config.Config) *Agent {
	return &Agent{
		cfg:    cfg,
		orch:   client.New(cfg.OrchestratorURL, cfg.NodeHMACSecret),
		status: "Ready",
	}
}

func (a *Agent) setJob(job *client.ClaimedJob) {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.currentJob = job
	if job != nil {
		a.status = "Busy"
	} else {
		a.status = "Ready"
	}
}

func (a *Agent) getStatus() (string, *string) {
	a.mu.Lock()
	defer a.mu.Unlock()
	if a.currentJob != nil {
		jobID := a.currentJob.JobID
		return a.status, &jobID
	}
	return a.status, nil
}

// ── Main ──────────────────────────────────────────────────────────────────────

func main() {
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)
	log.Printf("FlatRender Node Agent v%s starting (OS: %s, Arch: %s)",
		"0.1.0", runtime.GOOS, runtime.GOARCH)

	cfg, err := config.Load()
	if err != nil {
		log.Fatalf("config: %v", err)
	}
	log.Printf("Node ID: %s | Region: %q | Orchestrator: %s",
		cfg.NodeID, cfg.Region, cfg.OrchestratorURL)
	if cfg.AEPath == "" {
		log.Printf("AE_PATH not set — using mock renderer (development mode)")
	} else {
		log.Printf("AE binary: %s", cfg.AEPath)
	}

	agent := newAgent(cfg)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Graceful shutdown on SIGTERM / SIGINT
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGTERM, syscall.SIGINT)
	go func() {
		s := <-sigs
		log.Printf("received %s, shutting down…", s)
		cancel()
	}()

	// Register as online
	if err := agent.registerOnline(ctx); err != nil {
		log.Fatalf("failed to register with orchestrator: %v", err)
	}

	// Start health endpoint
	go agent.serveHealth(ctx)

	// Main loops
	var wg sync.WaitGroup
	wg.Add(2)
	go func() { defer wg.Done(); agent.heartbeatLoop(ctx) }()
	go func() { defer wg.Done(); agent.pollLoop(ctx) }()
	wg.Wait()
	log.Printf("shutdown complete")
}

// ── Registration ──────────────────────────────────────────────────────────────

func (a *Agent) registerOnline(ctx context.Context) error {
	req := client.OnlineRequest{
		NodeAgentVersion:    a.cfg.AgentVersion,
		CurrentAEVersion:    a.cfg.AEVersion,
		AvailableAEVersions: []string{a.cfg.AEVersion},
		CachedTemplateMD5s:  []string{},
	}
	if err := a.orch.Online(ctx, a.cfg.NodeID, req); err != nil {
		return fmt.Errorf("online: %w", err)
	}
	log.Printf("registered as online with orchestrator")
	return nil
}

// ── Heartbeat loop ────────────────────────────────────────────────────────────

func (a *Agent) heartbeatLoop(ctx context.Context) {
	ticker := time.NewTicker(time.Duration(a.cfg.HeartbeatIntervalSec) * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			a.sendHeartbeat(ctx)
		}
	}
}

func (a *Agent) sendHeartbeat(ctx context.Context) {
	status, jobID := a.getStatus()
	req := client.HeartbeatRequest{
		NodeID:       a.cfg.NodeID,
		Status:       status,
		CurrentJobID: jobID,
	}
	hbCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()

	resp, err := a.orch.Heartbeat(hbCtx, a.cfg.NodeID, req)
	if err != nil {
		log.Printf("heartbeat error: %v", err)
		return
	}
	if len(resp.PendingCommands) > 0 {
		log.Printf("orchestrator commands: %v", resp.PendingCommands)
	}
}

// ── Job poll loop ─────────────────────────────────────────────────────────────

func (a *Agent) pollLoop(ctx context.Context) {
	ticker := time.NewTicker(time.Duration(a.cfg.PollIntervalSec) * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			status, _ := a.getStatus()
			if status == "Busy" {
				continue // already rendering
			}
			a.tryClaimAndRun(ctx)
		}
	}
}

func (a *Agent) tryClaimAndRun(ctx context.Context) {
	claimCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()

	job, err := a.orch.ClaimJob(claimCtx, a.cfg.NodeID, a.cfg.Region)
	if err != nil {
		log.Printf("claim error: %v", err)
		return
	}
	if job == nil {
		return // queue empty
	}

	log.Printf("claimed job %s (project %s, %s %s %dfps)",
		job.JobID, job.SavedProjectID, job.Quality, job.Resolution, job.FrameRate)

	a.setJob(job)
	go func() {
		defer a.setJob(nil)
		a.runJob(ctx, job)
	}()
}

// ── Render execution ──────────────────────────────────────────────────────────

func (a *Agent) runJob(ctx context.Context, job *client.ClaimedJob) {
	log.Printf("[job %s] starting render", job.JobID)

	// ── Step 1: Download .aep template ───────────────────────────────────────
	aepPath := ""
	if job.AEPDownloadURL != "" && a.cfg.AEPath != "" {
		localAEP := filepath.Join(a.cfg.WorkDir, "templates", job.JobID, "template.aep")
		dlCtx, dlCancel := context.WithTimeout(ctx, 10*time.Minute)
		n, dlErr := runner.DownloadFile(dlCtx, job.AEPDownloadURL, localAEP)
		dlCancel()
		if dlErr != nil {
			log.Printf("[job %s] AEP download failed (%v) — falling back to mock", job.JobID, dlErr)
		} else {
			log.Printf("[job %s] AEP downloaded (%d bytes) → %s", job.JobID, n, localAEP)
			aepPath = localAEP
		}
	}

	rJob := &runner.Job{
		JobID:          job.JobID,
		SavedProjectID: job.SavedProjectID,
		Quality:        job.Quality,
		Resolution:     job.Resolution,
		FrameRate:      job.FrameRate,
		HasMusic:       job.HasMusic,
		HasVoiceover:   job.HasVoiceover,
		AEPFilePath:    aepPath,
	}

	onProgress := func(ctx context.Context, pct int, msg string) error {
		log.Printf("[job %s] %d%% %s", job.JobID, pct, msg)
		return nil
	}

	onPreview := func(ctx context.Context, imageB64 string) error {
		pvCtx, cancel := context.WithTimeout(ctx, 8*time.Second)
		defer cancel()
		if err := a.orch.UpdatePreview(pvCtx, job.JobID, imageB64); err != nil {
			log.Printf("[job %s] preview push error: %v", job.JobID, err)
		}
		return nil
	}

	// ── Step 2: Render ───────────────────────────────────────────────────────
	outputPath, err := runner.Run(ctx, a.cfg.AEPath, a.cfg.WorkDir, rJob, onProgress, onPreview)
	if err != nil {
		if ctx.Err() != nil {
			log.Printf("[job %s] render cancelled", job.JobID)
			return
		}
		log.Printf("[job %s] render failed: %v", job.JobID, err)
		failCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		if ferr := a.orch.Fail(failCtx, job.JobID, err.Error(), "Rendering"); ferr != nil {
			log.Printf("[job %s] fail report error: %v", job.JobID, ferr)
		}
		return
	}
	log.Printf("[job %s] render done → %s", job.JobID, outputPath)

	// ── Step 3: Get presigned upload URL + upload output to MinIO ─────────────
	var exportID *string
	uploadCtx, uploadCancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer uploadCancel()

	uploadInfo, urlErr := a.orch.GetOutputUploadURL(uploadCtx, job.JobID)
	if urlErr != nil {
		log.Printf("[job %s] get upload URL failed: %v — completing without export", job.JobID, urlErr)
	} else {
		log.Printf("[job %s] uploading output to %s", job.JobID, uploadInfo.ObjectKey)
		if _, upErr := runner.UploadFile(uploadCtx, uploadInfo.UploadURL, outputPath); upErr != nil {
			log.Printf("[job %s] upload failed: %v — completing without export", job.JobID, upErr)
		} else {
			log.Printf("[job %s] upload complete (export %s)", job.JobID, uploadInfo.ExportID)
			exportID = &uploadInfo.ExportID
		}
	}

	// ── Step 4: Report complete ───────────────────────────────────────────────
	completeCtx, completeCancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer completeCancel()
	if err := a.orch.Complete(completeCtx, job.JobID, exportID); err != nil {
		log.Printf("[job %s] complete report error: %v", job.JobID, err)
	} else {
		log.Printf("[job %s] reported as completed (export=%v)", job.JobID, exportID)
	}
}

// ── Health endpoint ───────────────────────────────────────────────────────────

func (a *Agent) serveHealth(ctx context.Context) {
	mux := http.NewServeMux()
	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		status, jobID := a.getStatus()
		w.Header().Set("Content-Type", "application/json")
		resp := map[string]any{
			"ok":         true,
			"node_id":    a.cfg.NodeID,
			"status":     status,
			"current_job": jobID,
			"version":    a.cfg.AgentVersion,
		}
		_ = json.NewEncoder(w).Encode(resp)
	})

	srv := &http.Server{
		Addr:    fmt.Sprintf(":%d", a.cfg.ListenPort),
		Handler: mux,
	}

	go func() {
		<-ctx.Done()
		_ = srv.Shutdown(context.Background())
	}()

	log.Printf("health endpoint listening on :%d", a.cfg.ListenPort)
	if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		log.Printf("health server error: %v", err)
	}
}