Refactor imports and clean up whitespace in API and runner client files. Update auth package references to use the new package structure. Improve mutex usage in the runner client for better concurrency handling.

This commit is contained in:
2025-11-22 05:45:13 -06:00
parent fb2e318eaa
commit 27a09aedd6
3 changed files with 31 additions and 34 deletions

View File

@@ -7,8 +7,6 @@ import (
"net/http" "net/http"
"time" "time"
"github.com/go-chi/chi/v5"
"fuego/internal/auth"
"fuego/pkg/types" "fuego/pkg/types"
) )
@@ -22,7 +20,7 @@ func (s *Server) handleGenerateRegistrationToken(w http.ResponseWriter, r *http.
// Default expiration: 24 hours // Default expiration: 24 hours
expiresIn := 24 * time.Hour expiresIn := 24 * time.Hour
var req struct { var req struct {
ExpiresInHours int `json:"expires_in_hours,omitempty"` ExpiresInHours int `json:"expires_in_hours,omitempty"`
} }
@@ -39,9 +37,9 @@ func (s *Server) handleGenerateRegistrationToken(w http.ResponseWriter, r *http.
} }
s.respondJSON(w, http.StatusCreated, map[string]interface{}{ s.respondJSON(w, http.StatusCreated, map[string]interface{}{
"token": token, "token": token,
"expires_in": expiresIn.String(), "expires_in": expiresIn.String(),
"expires_at": time.Now().Add(expiresIn), "expires_at": time.Now().Add(expiresIn),
}) })
} }
@@ -169,4 +167,3 @@ func (s *Server) handleListRunnersAdmin(w http.ResponseWriter, r *http.Request)
s.respondJSON(w, http.StatusOK, runners) s.respondJSON(w, http.StatusOK, runners)
} }

View File

@@ -10,35 +10,36 @@ import (
"sync" "sync"
"time" "time"
authpkg "fuego/internal/auth"
"fuego/internal/database"
"fuego/internal/storage"
"fuego/pkg/types"
"github.com/go-chi/chi/v5" "github.com/go-chi/chi/v5"
"github.com/go-chi/chi/v5/middleware" "github.com/go-chi/chi/v5/middleware"
"github.com/go-chi/cors" "github.com/go-chi/cors"
"github.com/gorilla/websocket" "github.com/gorilla/websocket"
"fuego/internal/auth"
"fuego/internal/database"
"fuego/internal/storage"
"fuego/pkg/types"
) )
// Server represents the API server // Server represents the API server
type Server struct { type Server struct {
db *database.DB db *database.DB
auth *auth.Auth auth *authpkg.Auth
secrets *auth.Secrets secrets *authpkg.Secrets
storage *storage.Storage storage *storage.Storage
router *chi.Mux router *chi.Mux
// WebSocket connections // WebSocket connections
wsUpgrader websocket.Upgrader wsUpgrader websocket.Upgrader
runnerConns map[int64]*websocket.Conn runnerConns map[int64]*websocket.Conn
runnerConnsMu sync.RWMutex runnerConnsMu sync.RWMutex
frontendConns map[string]*websocket.Conn // key: "jobId:taskId" frontendConns map[string]*websocket.Conn // key: "jobId:taskId"
frontendConnsMu sync.RWMutex frontendConnsMu sync.RWMutex
} }
// NewServer creates a new API server // NewServer creates a new API server
func NewServer(db *database.DB, auth *auth.Auth, storage *storage.Storage) (*Server, error) { func NewServer(db *database.DB, auth *authpkg.Auth, storage *storage.Storage) (*Server, error) {
secrets, err := auth.NewSecrets(db.DB) secrets, err := authpkg.NewSecrets(db.DB)
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to initialize secrets: %w", err) return nil, fmt.Errorf("failed to initialize secrets: %w", err)
} }
@@ -142,10 +143,10 @@ func (s *Server) setupRoutes() {
s.router.Route("/api/runner", func(r chi.Router) { s.router.Route("/api/runner", func(r chi.Router) {
// Registration doesn't require auth (uses token) // Registration doesn't require auth (uses token)
r.Post("/register", s.handleRegisterRunner) r.Post("/register", s.handleRegisterRunner)
// WebSocket endpoint (auth handled in handler) // WebSocket endpoint (auth handled in handler)
r.Get("/ws", s.handleRunnerWebSocket) r.Get("/ws", s.handleRunnerWebSocket)
// File operations still use HTTP (WebSocket not suitable for large files) // File operations still use HTTP (WebSocket not suitable for large files)
r.Group(func(r chi.Router) { r.Group(func(r chi.Router) {
r.Use(func(next http.Handler) http.Handler { r.Use(func(next http.Handler) http.Handler {
@@ -290,7 +291,7 @@ func (s *Server) handleGetMe(w http.ResponseWriter, r *http.Request) {
// Helper to get user ID from context // Helper to get user ID from context
func getUserID(r *http.Request) (int64, error) { func getUserID(r *http.Request) (int64, error) {
userID, ok := auth.GetUserID(r.Context()) userID, ok := authpkg.GetUserID(r.Context())
if !ok { if !ok {
return 0, fmt.Errorf("user ID not found in context") return 0, fmt.Errorf("user ID not found in context")
} }
@@ -436,7 +437,7 @@ func (s *Server) recoverStuckTasks() {
// Check for task timeouts // Check for task timeouts
s.recoverTaskTimeouts() s.recoverTaskTimeouts()
// Distribute newly recovered tasks // Distribute newly recovered tasks
s.distributeTasksToRunners() s.distributeTasksToRunners()
}() }()
@@ -510,4 +511,3 @@ func (s *Server) recoverTaskTimeouts() {
} }
} }
} }

View File

@@ -20,8 +20,9 @@ import (
"sync" "sync"
"time" "time"
"github.com/gorilla/websocket"
"fuego/pkg/types" "fuego/pkg/types"
"github.com/gorilla/websocket"
) )
// Client represents a runner client // Client represents a runner client
@@ -35,7 +36,7 @@ type Client struct {
runnerSecret string runnerSecret string
managerSecret string managerSecret string
wsConn *websocket.Conn wsConn *websocket.Conn
wsConnMu sync.Mutex wsConnMu sync.RWMutex
stopChan chan struct{} stopChan chan struct{}
} }
@@ -144,7 +145,7 @@ func (c *Client) ConnectWebSocket() error {
h := hmac.New(sha256.New, []byte(c.runnerSecret)) h := hmac.New(sha256.New, []byte(c.runnerSecret))
h.Write([]byte(message)) h.Write([]byte(message))
signature := hex.EncodeToString(h.Sum(nil)) signature := hex.EncodeToString(h.Sum(nil))
// Convert HTTP URL to WebSocket URL // Convert HTTP URL to WebSocket URL
wsURL := strings.Replace(c.managerURL, "http://", "ws://", 1) wsURL := strings.Replace(c.managerURL, "http://", "ws://", 1)
wsURL = strings.Replace(wsURL, "https://", "wss://", 1) wsURL = strings.Replace(wsURL, "https://", "wss://", 1)
@@ -268,8 +269,8 @@ func (c *Client) handleTaskAssignment(msg map[string]interface{}) {
// Convert to task map format // Convert to task map format
taskMap := map[string]interface{}{ taskMap := map[string]interface{}{
"id": taskID, "id": taskID,
"job_id": jobID, "job_id": jobID,
"frame_start": frameStart, "frame_start": frameStart,
"frame_end": frameEnd, "frame_end": frameEnd,
} }
@@ -402,7 +403,7 @@ func (c *Client) processTask(task map[string]interface{}, jobName, outputFormat
// Step: render_blender // Step: render_blender
c.sendStepUpdate(taskID, "render_blender", types.StepStatusRunning, "") c.sendStepUpdate(taskID, "render_blender", types.StepStatusRunning, "")
c.sendLog(taskID, types.LogLevelInfo, fmt.Sprintf("Starting Blender render for frame %d...", frameStart), "render_blender") c.sendLog(taskID, types.LogLevelInfo, fmt.Sprintf("Starting Blender render for frame %d...", frameStart), "render_blender")
// Execute Blender // Execute Blender
cmd := exec.Command("blender", "-b", blendFile, "-o", outputPattern, "-f", fmt.Sprintf("%d", frameStart)) cmd := exec.Command("blender", "-b", blendFile, "-o", outputPattern, "-f", fmt.Sprintf("%d", frameStart))
cmd.Dir = workDir cmd.Dir = workDir
@@ -432,7 +433,7 @@ func (c *Client) processTask(task map[string]interface{}, jobName, outputFormat
} }
c.sendStepUpdate(taskID, uploadStepName, types.StepStatusRunning, "") c.sendStepUpdate(taskID, uploadStepName, types.StepStatusRunning, "")
c.sendLog(taskID, types.LogLevelInfo, "Uploading output file...", uploadStepName) c.sendLog(taskID, types.LogLevelInfo, "Uploading output file...", uploadStepName)
outputPath, err := c.uploadFile(jobID, outputFile) outputPath, err := c.uploadFile(jobID, outputFile)
if err != nil { if err != nil {
errMsg := fmt.Sprintf("failed to upload output: %w", err) errMsg := fmt.Sprintf("failed to upload output: %w", err)
@@ -446,7 +447,7 @@ func (c *Client) processTask(task map[string]interface{}, jobName, outputFormat
// Step: complete // Step: complete
c.sendStepUpdate(taskID, "complete", types.StepStatusRunning, "") c.sendStepUpdate(taskID, "complete", types.StepStatusRunning, "")
c.sendLog(taskID, types.LogLevelInfo, "Task completed successfully", "complete") c.sendLog(taskID, types.LogLevelInfo, "Task completed successfully", "complete")
// Mark task as complete // Mark task as complete
if err := c.completeTask(taskID, outputPath, true, ""); err != nil { if err := c.completeTask(taskID, outputPath, true, ""); err != nil {
c.sendStepUpdate(taskID, "complete", types.StepStatusFailed, err.Error()) c.sendStepUpdate(taskID, "complete", types.StepStatusFailed, err.Error())
@@ -731,7 +732,7 @@ func (c *Client) uploadFile(jobID int64, filePath string) (string, error) {
// Create multipart form // Create multipart form
var buf bytes.Buffer var buf bytes.Buffer
formWriter := multipart.NewWriter(&buf) formWriter := multipart.NewWriter(&buf)
part, err := formWriter.CreateFormFile("file", filepath.Base(filePath)) part, err := formWriter.CreateFormFile("file", filepath.Base(filePath))
if err != nil { if err != nil {
return "", fmt.Errorf("failed to create form file: %w", err) return "", fmt.Errorf("failed to create form file: %w", err)
@@ -813,4 +814,3 @@ func (c *Client) sendTaskComplete(taskID int64, outputPath string, success bool,
} }
return fmt.Errorf("WebSocket not connected, cannot complete task") return fmt.Errorf("WebSocket not connected, cannot complete task")
} }